#include "log-fourier.hpp"
#include "p-spin.hpp"
#include <getopt.h>
#include <iostream>

int main(int argc, char* argv[]) {
  unsigned p = 2;
  unsigned s = 2;
  Real λ = 0.5;
  Real τ₀ = 0;

  unsigned log2n = 8;
  Real Δτ = 0.1;
  Real k = 0.1;

  Real ε = 1e-13;
  Real γ = 1;
  Real βₘₐₓ = 0.7;
  Real Δβ = 0.01;

  int opt;

  while ((opt = getopt(argc, argv, "p:s:2:T:t:b:d:g:k:D:")) != -1) {
    switch (opt) {
    case 'p':
      p = atoi(optarg);
      break;
    case 's':
      s = atoi(optarg);
      break;
    case '2':
      log2n = atoi(optarg);
      break;
    case 't':
      τ₀ = atof(optarg);
      break;
    case 'b':
      βₘₐₓ = atof(optarg);
      break;
    case 'd':
      Δβ = atof(optarg);
      break;
    case 'g':
      γ = atof(optarg);
      break;
    case 'k':
      k = atof(optarg);
      break;
    case 'D':
      Δτ = atof(optarg);
      break;
    default:
      exit(1);
    }
  }

  unsigned N = pow(2, log2n);

  LogarithmicFourierTransform fft(N, k, Δτ, 4);

  Real Γ₀ = 1.0;
  Real μ = Γ₀;
  if (τ₀ > 0) {
    μ = (sqrt(1+4*Γ₀*τ₀)-1)/(2*τ₀);
  }

  std::vector<Real> C(N);
  std::vector<Real> R(N);
  std::vector<Complex> Ct(N);
  std::vector<Complex> Rt(N);

  // start from the exact solution for β = 0
  for (unsigned n = 0; n < N; n++) {
    if (τ₀ > 0) {
      C[n] = Γ₀ * (exp(-μ * fft.t(n)) - μ * τ₀ * exp(-fft.t(n) / τ₀)) / (μ - pow(μ, 3) * pow(τ₀, 2));
    } else {
      C[n] = Γ₀ * exp(-μ * fft.t(n)) / μ;
    }
    R[n] = exp(-μ * fft.t(n));

    Ct[n] = 2 * Γ₀ / (pow(μ, 2) + pow(fft.ν(n), 2)) / (1 + pow(τ₀ * fft.ν(n), 2));
    Rt[n] = 1.0 / (μ + 1i * fft.ν(n));
  }

  Real β = 0;
  while (β < βₘₐₓ) {
  Real ΔC = 100;
  while (ΔC > ε) {
    std::vector<Real> RddfC(N);
    std::vector<Real> dfC(N);
    for (unsigned n = 0; n < N; n++) {
      RddfC[n] = R[n] * ddf(λ, p, s, C[n]);
      dfC[n] = df(λ, p, s, C[n]);
    }
    std::vector<Complex> RddfCt = fft.fourier(RddfC, false);
    std::vector<Complex> dfCt = fft.fourier(dfC, true);

    std::vector<Complex> Rtnew(N);
    std::vector<Complex> Ctnew(N);

    for (unsigned n = 0; n < N; n++) {
      Rtnew[n] = (1.0 + pow(β, 2) * RddfCt[n] * Rt[n]) / (μ + 1i * fft.ν(n));
      Ctnew[n] = (2 * Γ₀ * std::conj(Rtnew[n]) / (1 + pow(τ₀ * fft.ν(n), 2)) + pow(β, 2) * (RddfCt[n] * Ct[n] + dfCt[n] * std::conj(Rtnew[n]))) / (μ + 1i * fft.ν(n));
//      Ctnew[n] = - 2 * Γ₀ * Rtnew[n].imag() / (1 + pow(τ₀ * fft.ν(n), 2)) / fft.ν(n);
    }

    std::vector<Real> Cnew = fft.inverse(Ctnew);
    std::vector<Real> Rnew = fft.inverse(Rtnew);

    ΔC = 0;
    for (unsigned i = 0; i < N; i++) {
      ΔC += std::norm(Ct[i] - Ctnew[i]);
      ΔC += std::norm(Rt[i] - Rtnew[i]);
    }
    ΔC = sqrt(ΔC) / (2*N);

    for (unsigned i = 0; i < N; i++) {
      C[i] += γ * (Cnew[i] - C[i]);
      R[i] += γ * (Rnew[i] - R[i]);
      Ct[i] += γ * (Ctnew[i] - Ct[i]);
      Rt[i] += γ * (Rtnew[i] - Rt[i]);
    }

    μ *= C[0];

//    std::cerr << ΔC << std::endl;
  }

  /* Integrate the energy using Simpson's rule */
  Real E = 0;
  for (unsigned i = 0; i < N/2-1; i++) {
    Real h₂ᵢ   = fft.t(2*i+1) - fft.t(2*i);
    Real h₂ᵢ₊₁ = fft.t(2*i+2) - fft.t(2*i+1);
    Real f₂ᵢ   = R[2*i]   * df(λ, p, s, C[2*i]);
    Real f₂ᵢ₊₁ = R[2*i+1] * df(λ, p, s, C[2*i+1]);
    Real f₂ᵢ₊₂ = R[2*i+2] * df(λ, p, s, C[2*i+2]);
    E += (h₂ᵢ + h₂ᵢ₊₁) / 6 * (
          (2 - h₂ᵢ₊₁ / h₂ᵢ) * f₂ᵢ
          + pow(h₂ᵢ + h₂ᵢ₊₁, 2) / (h₂ᵢ * h₂ᵢ₊₁) * f₂ᵢ₊₁
          + (2 - h₂ᵢ / h₂ᵢ₊₁) * f₂ᵢ₊₂
        );
  }
  E *= β;

  std::cerr << β << " " << μ << " " << Ct[0].real() << " " << E << std::endl;
   β += Δβ;
  }

  for (unsigned i = 0; i < N; i++) {
    std::cout << fft.t(i) << " " << C[i] << std::endl;
  }

  return 0;
}