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

std::string logFourierFile(std::string prefix, unsigned p, unsigned s, Real λ, Real τ₀, Real β, unsigned log2n, Real Δτ, Real k) {
  return prefix + "_" + std::to_string(p) + "_" + std::to_string(s) + "_" + std::to_string(λ) + "_" + std::to_string(τ₀) + "_" + std::to_string(β) + "_" + std::to_string(log2n) + "_" + std::to_string(Δτ)  + "_" + std::to_string(k) + ".dat";
}

std::tuple<std::vector<Complex>, std::vector<Complex>> RddfCtdfCt(LogarithmicFourierTransform& fft, const std::vector<Real>& C, const std::vector<Real>& R, unsigned p, unsigned s, Real λ) {
  std::vector<Real> dfC(C.size());
  std::vector<Real> RddfC(C.size());
  for (unsigned n = 0; n < C.size(); 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);

  return {RddfCt, dfCt};
}

Real estimateZ(LogarithmicFourierTransform& fft, const std::vector<Real>& C, const std::vector<Complex>& Ct, const std::vector<Real>& R, const std::vector<Complex>& Rt, unsigned p, unsigned s, Real λ, Real τ₀, Real β) {
  auto [RddfCt, dfCt] = RddfCtdfCt(fft, C, R, p, s, λ);
  Real Γ₀ = 1 + τ₀ / 2;

  return ((2 * Γ₀ * std::conj(Rt[0]) + pow(β, 2) * (RddfCt[0] * Ct[0] + dfCt[0] * std::conj(Rt[0]))) / Ct[0]).real();
}

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

  /* Log-Fourier parameters */
  unsigned log2n = 8;
  Real Δτ = 0.1;
  Real k = 0.1;

  /* Iteration parameters */
  Real ε = 1e-14;
  Real γ = 1;
  Real β₀ = 0;
  Real βₘₐₓ = 0.7;
  Real Δβ = 0.01;
  bool loadData = false;

  int opt;

  while ((opt = getopt(argc, argv, "p:s:2:T:t:b:d:g:k:D:e:0:l")) != -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;
    case 'e':
      ε = atof(optarg);
      break;
    case '0':
      β₀ = atof(optarg);
      break;
    case 'l':
      loadData = true;
      break;
    default:
      exit(1);
    }
  }

  unsigned N = pow(2, log2n);

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

  Real Γ₀ = 1.0 + τ₀;
  Real μₜ₋₁ = 1.0;

  std::vector<Real> Cₜ₋₁(N);
  std::vector<Real> Rₜ₋₁(N);
  std::vector<Complex> Ĉₜ₋₁(N);
  std::vector<Complex> Ȓₜ₋₁(N);

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

      Ĉₜ₋₁[n] = 2 * Γ₀ / (pow(μₜ₋₁, 2) + pow(fft.ν(n), 2)) / (1 + pow(τ₀ * fft.ν(n), 2));
      Ȓₜ₋₁[n] = 1.0 / (μₜ₋₁ + 1i * fft.ν(n));
    }
  } else {
    std::ifstream cfile(logFourierFile("C", p, s, λ, τ₀, β₀, log2n, Δτ, k), std::ios::binary);
    cfile.read((char*)(Cₜ₋₁.data()), N * sizeof(Real));
    cfile.close();
    std::ifstream rfile(logFourierFile("R", p, s, λ, τ₀, β₀, log2n, Δτ, k), std::ios::binary);
    rfile.read((char*)(Rₜ₋₁.data()), N * sizeof(Real));
    rfile.close();

    Ĉₜ₋₁ = fft.fourier(Cₜ₋₁, true);
    Ȓₜ₋₁ = fft.fourier(Rₜ₋₁, false);

    μₜ₋₁ = estimateZ(fft, Cₜ₋₁, Ĉₜ₋₁, Rₜ₋₁, Ȓₜ₋₁, p, s, λ, τ₀, β₀);
  }

  std::vector<Real> Cₜ = Cₜ₋₁;
  std::vector<Real> Rₜ = Rₜ₋₁;
  std::vector<Complex> Ĉₜ = Ĉₜ₋₁;
  std::vector<Complex> Ȓₜ = Ȓₜ₋₁;
  Real μₜ = μₜ₋₁;

  Real β = β₀ + Δβ;
  while (β < βₘₐₓ) {
    Real ΔC = 100;
    while (ΔC > ε) {
      auto [RddfCt, dfCt] = RddfCtdfCt(fft, Cₜ, Rₜ, p, s, λ);

      std::vector<Complex> Ĉₜ₊₁(N);
      std::vector<Complex> Ȓₜ₊₁(N);
      for (unsigned n = 0; n < N; n++) {
        Ȓₜ₊₁[n] = (1.0 + pow(β, 2) * RddfCt[n] * Ȓₜ[n]) / (μₜ + 1i * fft.ν(n));
        Ĉₜ₊₁[n] = (2 * Γ₀ * std::conj(Ȓₜ[n]) / (1 + pow(τ₀ * fft.ν(n), 2)) + pow(β, 2) * (RddfCt[n] * Ĉₜ[n] + dfCt[n] * std::conj(Ȓₜ[n]))) / (μₜ + 1i * fft.ν(n));
      }
      std::vector<Real> Rₜ₊₁ = fft.inverse(Ȓₜ₊₁);
      std::vector<Real> Cₜ₊₁ = fft.inverse(Ĉₜ₊₁);

      μₜ *= pow(tanh(Cₜ₊₁[0]-1)+1, 0.05);

      ΔC = 0;
      for (unsigned i = 0; i < N; i++) {
        ΔC += std::norm(Ĉₜ[i] - Ĉₜ₊₁[i]);
        ΔC += std::norm(Ȓₜ[i] - Ȓₜ₊₁[i]);
      }
      ΔC = sqrt(ΔC) / (2*N);

      for (unsigned i = 0; i < N; i++) {
        Cₜ[i] += γ * (Cₜ₊₁[i] - Cₜ[i]);
        Rₜ[i] += γ * (Rₜ₊₁[i] - Rₜ[i]);
        Ĉₜ[i] += γ * (Ĉₜ₊₁[i] - Ĉₜ[i]);
        Ȓₜ[i] += γ * (Ȓₜ₊₁[i] - Ȓₜ[i]);
      }

      std::cerr << "\x1b[2K" << "\r";
      std::cerr << β << " " << μₜ << " " << ΔC << " " << γ << " " << Cₜ[0];
    }

    if (std::isnan(Cₜ[0])) {
      Cₜ = Cₜ₋₁;
      Rₜ = Rₜ₋₁;
      Ĉₜ = Ĉₜ₋₁;
      Ȓₜ = Ȓₜ₋₁;
      μₜ = μₜ₋₁;
      γ /= 2;
    } else {
      /* Integrate the energy using Simpson's rule */
      Real E = 0;
      for (unsigned n = 0; n < N/2-1; n++) {
        Real h₂ₙ   = fft.t(2*n+1) - fft.t(2*n);
        Real h₂ₙ₊₁ = fft.t(2*n+2) - fft.t(2*n+1);
        Real f₂ₙ   = Rₜ[2*n]   * df(λ, p, s, Cₜ[2*n]);
        Real f₂ₙ₊₁ = Rₜ[2*n+1] * df(λ, p, s, Cₜ[2*n+1]);
        Real f₂ₙ₊₂ = Rₜ[2*n+2] * df(λ, p, s, Cₜ[2*n+2]);
        E += (h₂ₙ + h₂ₙ₊₁) / 6 * (
              (2 - h₂ₙ₊₁ / h₂ₙ) * f₂ₙ
              + pow(h₂ₙ + h₂ₙ₊₁, 2) / (h₂ₙ * h₂ₙ₊₁) * f₂ₙ₊₁
              + (2 - h₂ₙ / h₂ₙ₊₁) * f₂ₙ₊₂
            );
      }
      E *= β;

      std::cerr << "\x1b[2K" << "\r";
      std::cerr << β << " " << μₜ << " " << Ĉₜ[0].real() << " " << E << " " << γ << std::endl;

      std::ofstream outfile(logFourierFile("C", p, s, λ, τ₀, β, log2n, Δτ, k), std::ios::out | std::ios::binary);
      outfile.write((const char*)(Cₜ.data()), N * sizeof(Real));
      outfile.close();

      std::ofstream outfileR(logFourierFile("R", p, s, λ, τ₀, β, log2n, Δτ, k), std::ios::out | std::ios::binary);
      outfileR.write((const char*)(Rₜ.data()), N * sizeof(Real));
      outfileR.close();

      β += Δβ;
      Cₜ₋₁ = Cₜ;
      Rₜ₋₁ = Rₜ;
      Ĉₜ₋₁ = Ĉₜ;
      Ȓₜ₋₁ = Ȓₜ;
      μₜ₋₁ = μₜ;
    }
  }

  return 0;
}