#include <math.h>
#include <vector>
#include <complex>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <Eigen/Eigenvalues>

using namespace Eigen;

/* ***** ******************** ***** */
/* **** *//** Parameters **//* **** */
/* ***** ******************** ***** */

size_t latticeSize = 0;
bool doFerromagnet;
double Jxy;
double Jz;


/* **** ********************** **** */
/* *** *//** Declarations **//* *** */
/* **** ********************** **** */

struct magneticSubspace
{
    size_t m_size = 0;
    std::vector< size_t > m_binaryForm;
    size_t m_findState(size_t state)
    {
        size_t minState = 0, maxState = m_size - 1, resultState = 0;
        if((state < m_binaryForm[minState]) || (state > m_binaryForm[maxState]))
        {
            std::cout << "Binary Search Notification : state does not belong to given subspace." << std::endl;
            resultState = -1;
        }
        else while (true)
        {
            resultState = minState + (1 + maxState - minState) / 2;
            if (state > m_binaryForm[resultState]) minState = resultState + 1;
            else if (state < m_binaryForm[resultState]) maxState = resultState - 1;
            else if (minState > maxState)
            {
                std::cout << "Binary Search Notification : state does not belong to given subspace." << std::endl;
                resultState = -1;
                break;
            }
            else break;
        }
        return resultState;
    }
};

bool spinsAlingedFerro(size_t state, size_t siteI, size_t siteJ);
bool readInParameters(size_t & latticeSize, bool & doFerromagnet, double & Jxy, double & Jz);
bool readInParameters(double & delta, double & dk, double & dw, double & omegaMin, double & omegaMax);

size_t getNumberOfSpinsUp(size_t state);
size_t spinSwitch(size_t state, size_t siteI, size_t siteJ);
size_t spinFlipUp(size_t state, size_t siteI);

MatrixXd constructSubspaceHamiltonian(magneticSubspace & subspace);

void setDimensions(std::vector< magneticSubspace > & basis);
void constructBasis(size_t numberOfStates, std::vector< magneticSubspace > & basis);
void findGround(MatrixXd & subspaceHamiltonian, VectorXcd & groundState, double & groundEnergy);
void doSkPlus(VectorXcd & groundState, size_t ik, double k, magneticSubspace & groundSubspace, magneticSubspace & upSubspace, MatrixXcd & stateSkPlus);
void spectralFunction(MatrixXd & subspaceHamiltonian, magneticSubspace & groundSubspace, magneticSubspace & upSubspace, VectorXcd & groundState, double groundEnergy);

std::complex< double > I(0.,1.);
double PI = 3.14159265359;


/* **** ********************** **** */
/* *** *//** Main Program **//* *** */
/* **** ********************** **** */

int main()
{
    if(!readInParameters(latticeSize, doFerromagnet, Jxy, Jz)) return 1;

    std::cout << "Init : ";
        size_t groundStateSubspace = (latticeSize + 1) / 2;
        if (doFerromagnet)
        {
            Jxy = -Jxy;
            Jz = -Jz;
            groundStateSubspace = 0;
        }
        size_t numberOfStates = pow(2,latticeSize);
    std::cout << "Finished" << std::endl;

    std::cout << "Basis Construction : ";
        std::vector< magneticSubspace > basis(latticeSize + 1);
        setDimensions(basis);
        constructBasis(numberOfStates, basis);
    std::cout << "Finished" << std::endl;

    std::cout << "Ground State Searching : ";
        MatrixXd subspaceHamiltonian = constructSubspaceHamiltonian(basis[groundStateSubspace]);
        VectorXcd groundState = VectorXd::Zero(basis[groundStateSubspace].m_size);
        double groundEnergy;
        findGround(subspaceHamiltonian, groundState, groundEnergy);
    std::cout << "Finished" << std::endl;

    std::cout << "Spectral Function Calculation : ";
        subspaceHamiltonian = constructSubspaceHamiltonian(basis[groundStateSubspace + 1]);
        spectralFunction(subspaceHamiltonian, basis[groundStateSubspace], basis[groundStateSubspace + 1], groundState, groundEnergy);
    std::cout << "Finished" << std::endl << std::endl;

    std::cout << "Spectral function written to file : 'spectra.data'" << std::endl;
    std::cout << "Info written to file : 'spectra.info'" << std::endl << std::endl;

    system("PAUSE");
    return 0;
}


/* ***** ******************** ***** */
/* **** *//** Procedures **//* **** */
/* ***** ******************** ***** */

bool readInParameters(size_t & latticeSize, bool & doFerromagnet, double & Jxy, double & Jz)
{
    std::ifstream file;
    file.open("system.in");
    if(file.good())
    {
        file >> latticeSize;
        file >> doFerromagnet;
        file >> Jxy;
        file >> Jz;
        file.close();
    }
    if((latticeSize < 1) || ((Jxy == 0) && (Jz == 0))) return false;
    else return true;
}

size_t getNumberOfSpinsUp(size_t state)
{
    size_t numberOfSpinsUp = 0;
    for(size_t site = 0; site < latticeSize; ++site)
    {
        size_t spinUpAtSite = pow(2,site);
        if((state & spinUpAtSite) == spinUpAtSite) ++numberOfSpinsUp;
    }
    return numberOfSpinsUp;
}

void setDimensions(std::vector< magneticSubspace > & basis)
{
    std::vector< size_t > binomial(latticeSize + 1);
    binomial[0] = 1;
    for(size_t subspace = 1; subspace <= latticeSize; ++subspace)
    {
        binomial[subspace] = 1;
        for(size_t echo = subspace - 1; echo > 0; --echo)
            binomial[echo] += binomial[echo - 1];
    }
    for(size_t subspace = 0; subspace <= latticeSize; ++subspace)
        basis[subspace].m_binaryForm.resize(binomial[subspace]);
    return;
}

void constructBasis(size_t numberOfStates, std::vector< magneticSubspace > & basis)
{
    for(size_t state = 0; state < numberOfStates; ++state)
    {
        size_t subspace = getNumberOfSpinsUp(state);
        basis[subspace].m_binaryForm[basis[subspace].m_size] = state;
        basis[subspace].m_size += 1;
    }
    return;
}

size_t spinSwitch(size_t state, size_t siteI, size_t siteJ)
{
    siteI %= latticeSize;
    siteJ %= latticeSize;
    return state ^ size_t(pow(2,siteI) + pow(2,siteJ));
}

bool spinsAlingedFerro(size_t state, size_t siteI, size_t siteJ)
{
    bool spinI = state & size_t(pow(2,siteI));
    bool spinJ = state & size_t(pow(2,siteJ));
    return ((spinI && spinJ) || (!spinI && !spinJ));
}

MatrixXd constructSubspaceHamiltonian(magneticSubspace & subspace)
{
    MatrixXd subspaceHamiltonian = MatrixXd::Zero(subspace.m_size, subspace.m_size);
    for(size_t state = 0; state < subspace.m_size; ++state)
    {
        size_t binaryState = subspace.m_binaryForm[state];
        for(size_t siteI = 0; siteI < latticeSize; ++siteI)
        {
            size_t siteJ = (siteI + 1) % latticeSize;
            if(spinsAlingedFerro(binaryState,siteI,siteJ)) subspaceHamiltonian(state,state) += 0.25 * Jz;
            else
            {
                subspaceHamiltonian(state,state) -= 0.25 * Jz;
                subspaceHamiltonian(state,subspace.m_findState(spinSwitch(binaryState,siteI,siteJ))) += 0.5 * Jxy;
            }
        }
    }
    return subspaceHamiltonian;
}

void findGround(MatrixXd & subspaceHamiltonian, VectorXcd & groundState, double & groundEnergy)
{
    EigenSolver< MatrixXd > eigenSystem;
    eigenSystem.compute(subspaceHamiltonian, true);
    groundEnergy = eigenSystem.eigenvalues()[0].real();
    size_t statePosition = 0;
    size_t subspaceSize = size_t(groundState.size());
    for(size_t eigenState = 0; eigenState < subspaceSize; ++eigenState)
        if(eigenSystem.eigenvalues()[eigenState].real() < groundEnergy)
        {
            groundEnergy = eigenSystem.eigenvalues()[eigenState].real();
            statePosition = eigenState;
        };
    groundState = eigenSystem.eigenvectors().col(statePosition);
    return;
}

bool readInParameters(double & delta, double & dk, double & dw, double & omegaMin, double & omegaMax)
{
    std::ifstream file;
    file.open("green.in");
    if(file.good())
    {
        file >> delta;
        file >> dk;
        file >> dw;
        file >> omegaMin;
        file >> omegaMax;
        file.close();
    }
    if((delta < 0.0) || (dk < 0.0) || (dw < 0.0) || (omegaMin > omegaMax)) return false;
    else return true;
}

size_t spinFlipUp(size_t state, size_t siteI)
{
    size_t gate = size_t(pow(2, siteI % latticeSize));
    if(state & gate) return state;
    else return state ^ gate;
}

void doSkPlus(VectorXcd & groundState, size_t ik, double k, magneticSubspace & groundSubspace, magneticSubspace & upSubspace, MatrixXcd & stateSkPlus)
{
    for(size_t state = 0; state < groundSubspace.m_size; ++state)
    {
        size_t binaryState = groundSubspace.m_binaryForm[state];
        for(size_t site = 0; site < latticeSize; ++site)
        {
            double r = double(site);
            size_t spinFlipState = spinFlipUp(binaryState, site);
            if(spinFlipState != binaryState)
                stateSkPlus(upSubspace.m_findState(spinFlipState), ik) += groundState(state) * exp(I * k * r);
        }
    }
}

void spectralFunction(MatrixXd & subspaceHamiltonian, magneticSubspace & groundSubspace, magneticSubspace & upSubspace, VectorXcd & groundState, double groundEnergy)
{
    double delta, dk, dw, omegaMin, omegaMax;
    if(!readInParameters(delta, dk, dw, omegaMin, omegaMax))
    {
        std::cout << "Notification : Improper input in green.in file.";
        return;
    }
    EigenSolver< MatrixXd > eigenSystem;
    eigenSystem.compute(subspaceHamiltonian, true);

    size_t omegaPoints = size_t(ceil((omegaMax - omegaMin) / dw));

    VectorXd k(latticeSize + 1);
    VectorXd omega(omegaPoints);
    for(size_t ik = 0; ik <= latticeSize; ++ik) k(ik) = 2 * PI * ik / latticeSize;
    for(size_t it = 0; it < omegaPoints; ++it) omega(it) = omegaMin + it * (omegaMax - omegaMin) / omegaPoints;

    MatrixXcd denominator(omegaPoints, upSubspace.m_size);
    #pragma omp parallel
    {
        #pragma omp for
        for(size_t it = 0; it < omegaPoints; ++it)
            for(size_t state = 0; state < upSubspace.m_size; ++state)
                denominator(it,state) = omega(it) + I * delta - eigenSystem.eigenvalues()[state] + groundEnergy;
    }

    MatrixXcd stateSkPlus = MatrixXcd::Zero(upSubspace.m_size, latticeSize + 1);
    #pragma omp parallel
    {
        #pragma omp for
        for(size_t ik = 0; ik <= latticeSize; ++ik)
            doSkPlus(groundState, ik, k(ik), groundSubspace, upSubspace, stateSkPlus);
    }
    stateSkPlus.adjointInPlace();
    stateSkPlus = stateSkPlus * eigenSystem.eigenvectors();

    MatrixXcd greensFunction = MatrixXcd::Zero(omegaPoints, latticeSize + 1);
    #pragma omp parallel
    {
        #pragma omp for
        for(size_t ik = 0; ik <= latticeSize; ++ik)
        {
            for(size_t state = 0; state < upSubspace.m_size; ++state)
            {
                std::complex< double > numerator = pow(abs(stateSkPlus(ik,state)),2);
                for(size_t it = 0; it < omegaPoints; ++it) greensFunction(it,ik) += numerator / denominator(it,state);
            }
        }
    }

    MatrixXd spectra = MatrixXd::Zero(omegaPoints,latticeSize + 1);
    for(size_t it = 0; it < omegaPoints; ++it)
        for(size_t ik = 0; ik <= latticeSize; ++ik)
            spectra(it,ik) = -greensFunction(it,ik).imag() / PI;

    std::ofstream file;
    file.open("spectra.data");
    if(file.good())
    {
        file << spectra;
        file.close();
    }
    file.open("spectra.info");
    if(file.good())
    {
        file << latticeSize + 1 << std::endl << dk << std::endl << omegaPoints << std::endl << omegaMin << std::endl << omegaMax;
        file.close();
    }
    return;
}
