#include <getopt.h>
#include <vector>
#include <cmath>
#include <iostream>
#include <fftw3.h>
#include <complex>

using Real = double;
using Complex = std::complex<Real>;
using namespace std::complex_literals;

inline Real f(unsigned p, Real q) {
  return 0.5 * pow(q, p);
}

inline Real df(unsigned p, Real q) {
  return 0.5 * p * pow(q, p - 1);
}

inline Real ddf(unsigned p, Real q) {
  return 0.5 * p * (p - 1) * pow(q, p - 2);
}

class FourierTransform {
private:
  std::vector<Real> a;
  std::vector<Complex> â;
  fftw_plan plan_r2c;
  fftw_plan plan_c2r;
  Real Δω;
  Real Δτ;
public:
  FourierTransform(unsigned n, Real Δω, Real Δτ) : a(2 * n), â(n + 1), Δω(Δω), Δτ(Δτ) {
    plan_r2c = fftw_plan_dft_r2c_1d(2 * n, a.data(), reinterpret_cast<fftw_complex*>(â.data()), 0);
    plan_c2r = fftw_plan_dft_c2r_1d(2 * n, reinterpret_cast<fftw_complex*>(â.data()), a.data(), 0);
  }

  ~FourierTransform() {
    fftw_destroy_plan(plan_r2c);
    fftw_destroy_plan(plan_c2r);
    fftw_cleanup();
  }

  std::vector<Complex> fourier(const std::vector<Real>& c) {
    a = c;
    fftw_execute(plan_r2c);
    std::vector<Complex> ĉ(â.size());
    for (unsigned i = 0; i < â.size(); i++) {
      ĉ[i] = â[i] * (Δτ * M_PI);
    }
    return ĉ;
  }

  std::vector<Real> inverse(const std::vector<Complex>& ĉ) {
    â = ĉ;
    fftw_execute(plan_c2r);
    std::vector<Real> c(a.size());
    for (unsigned i = 0; i < a.size(); i++) {
      c[i] = a[i] * (Δω / (2 * M_PI));
    }
    return c;
  }
};

int main(int argc, char* argv[]) {
  unsigned p = 2;
  Real τ₀ = 0;
  Real yₘₐₓ = 0.5;
  Real Δy = 0.05;

  unsigned log2n = 8;
  Real τₘₐₓ = 20 / M_PI;

  unsigned maxIterations = 1000;
  Real ε = 1e-14;
  Real γ = 0;

  int opt;

  while ((opt = getopt(argc, argv, "p:2:T:t:y:d:I:g:")) != -1) {
    switch (opt) {
    case 'p':
      p = atoi(optarg);
      break;
    case '2':
      log2n = atoi(optarg);
      break;
    case 'T':
      τₘₐₓ = atof(optarg) / M_PI;
      break;
    case 't':
      τ₀ = atof(optarg);
      break;
    case 'y':
      yₘₐₓ = atof(optarg);
      break;
    case 'd':
      Δy = atof(optarg);
      break;
    case 'I':
      maxIterations = (unsigned)atof(optarg);
      break;
    case 'g':
      γ = atof(optarg);
      break;
    default:
      exit(1);
    }
  }

  unsigned n = pow(2, log2n);

  Real Δτ = τₘₐₓ / n;
  Real Δω = 1 / τₘₐₓ;

  Real z = (-1+sqrt(1+2*τ₀)) / (2 * τ₀);
  Real Γ₀ = 1;

  std::vector<Real> C(2 * n);
  std::vector<Real> R(2 * n);

  FourierTransform fft(n, Δω, Δτ);

  for (unsigned i = 0; i < n; i++) {
    Real τ = i * Δτ * M_PI;
    C[i] = Γ₀ / 2 * (exp(-z * τ) - z * τ₀ * exp(-τ / τ₀)) / (z - pow(z, 3) * pow(τ₀, 2));
    if (i > 0) {
      C[2 * n - i] = C[i];
    }
    R[i] = exp(-z * τ);
  }

  std::vector<Complex> Ct = fft.fourier(C);
  std::vector<Complex> Rt = fft.fourier(R);

  Real y = 0;

  while (y += Δy, y <= yₘₐₓ) {
    Real ΔC = 1;;
    unsigned it = 0;
    while (sqrt(ΔC / C.size()) > ε) {
      it++;
      std::vector<Real> RddfC(C.size());
      for (unsigned i = 0; i < C.size(); i++) {
        RddfC[i] = R[i] * ddf(p, C[i]);
      }
      std::vector<Complex> RddfCt = fft.fourier(RddfC);

      std::vector<Real> dfC(C.size());
      for (unsigned i = 0; i < C.size(); i++) {
        dfC[i] = df(p, C[i]);
      }
      std::vector<Complex> dfCt = fft.fourier(dfC);

      for (unsigned i = 0; i < Rt.size(); i++) {
        Real ω = i * Δω;
        Rt[i] = (1.0 + pow(y, 2) * RddfCt[i] * Rt[i]) / (z + 1i * ω);
      }

      for (unsigned i = 0; i < Ct.size(); i++) {
        Real ω = i * Δω;
        Ct[i] = (Γ₀ * std::conj(Rt[i]) / (1 + pow(τ₀ * ω, 2)) + pow(y, 2) * (RddfCt[i] * Ct[i] + dfCt[i] * std::conj(Rt[i]))) / (z + 1i * ω);
      }

      std::vector<Real> Cnew = fft.inverse(Ct);
      std::vector<Real> Rnew = fft.inverse(Rt);

      ΔC = 0;
      for (unsigned i = 0; i < Cnew.size(); i++) {
        ΔC += pow(Cnew[i] - C[i], 2);
        ΔC += pow(Rnew[i] - R[i], 2);
      }

      for (unsigned i = 0; i < Cnew.size(); i++) {
        C[i] += γ * (Cnew[i] - C[i]);
      }

      for (unsigned i = 0; i < Rnew.size(); i++) {
        R[i] += γ * (Rnew[i] - R[i]);
      }

      z *= Cnew[0];

      Real energy = 0;

      for (unsigned i = 0; i < n; i++) {
        energy += y * R[i] * df(p, C[i]) * M_PI * Δτ;
      }

      std::cerr << "y " << y << " iteration " << it << ": " << sqrt(ΔC / C.size()) << " " << pow(y, 2) / z << " " << energy << std::endl;

      if (it > maxIterations) {
        it = 0;
        γ /= 2;
      }
    }
  }

  for (unsigned i = 0; i < C.size(); i++) {
    std::cout << i * Δτ * M_PI << " " << C[i] << std::endl;
  }

  return 0;
}