#include <vector>
#include <random>
#include <cmath>
#include <fstream>
#include <iostream>
#include "ising.hpp"


IsingModel::IsingModel(int size, double temperature, bool randomInit, int seed)
    : latticeSize(size), temp(temperature), lattice(size, std::vector<int>(size, 1)), rng(seed) {
    // Precompute Boltzmann factors for efficiency
    boltzmannFactors.resize(5);
    for (int deltaE = -8; deltaE <= 8; deltaE += 4) {
        boltzmannFactors[(deltaE + 8) / 4] = exp(-deltaE / kB / temp);
        //std::cout << "Boltzmann factor for ΔE=" << deltaE << ": " << boltzmannFactors[(deltaE + 8) / 4] << "\n";
    }
    initializeLattice(randomInit);
}

void IsingModel::initializeLattice(bool randomInit) {
    if (!randomInit) {
        for (int i = 0; i < latticeSize; ++i) {
            for (int j = 0; j < latticeSize; ++j) {
                lattice[i][j] = 1; // All spins up
            }
        }
        return;
    }
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, 1);

    for (int i = 0; i < latticeSize; ++i) {
        for (int j = 0; j < latticeSize; ++j) {
            lattice[i][j] = dis(gen) == 0 ? -1 : 1; // Randomly assign -1 or +1
        }
    }
}

int IsingModel::getNeighborIndex(int i, int j, int dir, int offset) const {
    switch (dir) {
        case 0: return (i + offset) % latticeSize; // Down
        case 1: return (j + offset) % latticeSize; // Right
        case 2: return (i - offset + latticeSize) % latticeSize; // Up
        case 3: return (j - offset + latticeSize) % latticeSize; // Left
        default: return -1; // Invalid direction
    }
}
int IsingModel::getNeighborSum(int i, int j) const {
    return lattice[getNeighborIndex(i, j, 0)][j] + // Down
           lattice[i][getNeighborIndex(i, j, 1)] + // Right
           lattice[getNeighborIndex(i, j, 2)][j] + // Up
           lattice[i][getNeighborIndex(i, j, 3)];   // Left
}

double IsingModel::getBoltzmannFactor(int neighborSum, int spin) const {
    int deltaE = 2 * spin * neighborSum;
    return boltzmannFactors[(deltaE + 8) / 4];
}

void IsingModel::monteCarloStep() {
    std::uniform_int_distribution<> dist(0, latticeSize - 1);
    int i = dist(rng);
    int j = dist(rng);
    int neighborSum = getNeighborSum(i, j);
    int spin = lattice[i][j];
    double acceptanceProb = getBoltzmannFactor(neighborSum, spin);

    std::uniform_real_distribution<> probDist(0.0, 1.0);
    if (probDist(rng) < acceptanceProb) {
        lattice[i][j] *= -1; // Flip spin
    }
}

void IsingModel::monteCarloCycle() {
    for (int step = 0; step < latticeSize * latticeSize; ++step) {
        monteCarloStep();
    }
    updateAverages();
    recordEvolution();
}

void IsingModel::run(int cycles) {
    for (int cycle = 0; cycle < cycles; ++cycle) {
        monteCarloCycle();
    }
    //std::cout << "Accepted Moves: " << acceptedMoves << "\n";
    //std::cout << "Declined Moves: " << declinedMoves << "\n";
}

void IsingModel::updateAverages() {
    double mag = currentMagnetization(false);
    double en = currentEnergy(false);

    n_steps++;
    sumMagnetization += mag;
    sumSquaredMagnetization += mag * mag;
    sumEnergy += en;
    sumSquaredEnergy += en * en;
}

double IsingModel::currentMagnetization(bool normalized) const {
    double totalMagnetization = 0.0;
    for (int i = 0; i < latticeSize; ++i) {
        for (int j = 0; j < latticeSize; ++j) {
            totalMagnetization += lattice[i][j];
        }
    }
    return normalized? abs(totalMagnetization) / (latticeSize * latticeSize) : abs(totalMagnetization);
}

double IsingModel::currentEnergy(bool normalized) const {
    double totalEnergy = 0.0;
    for (int i = 0; i < latticeSize; ++i) {
        for (int j = 0; j < latticeSize; ++j) {
            int spin = lattice[i][j];
            int neighborSum = getNeighborSum(i, j);
            totalEnergy -= spin * neighborSum / 2.0; // Each pair counted twice
        }
    }
    return normalized ? totalEnergy / (latticeSize * latticeSize) : totalEnergy;
}

double IsingModel::calculateMagnetization(bool normalized) const {
    double avgMagnetization = sumMagnetization / n_steps;
    return normalized ? avgMagnetization / (latticeSize * latticeSize) : avgMagnetization;
}

double IsingModel::calculateEnergy(bool normalized) const {
    double avgEnergy = sumEnergy / n_steps;
    return normalized ? avgEnergy / (latticeSize * latticeSize) : avgEnergy;
}

double IsingModel::calculateMagnetizationVariance(bool normalized) const {
    double variance = (sumSquaredMagnetization / n_steps) - sumMagnetization * sumMagnetization / (n_steps * n_steps);
    return normalized ? variance / (latticeSize * latticeSize * latticeSize * latticeSize) : variance;
}

double IsingModel::calculateEnergyVariance(bool normalized) const {
    double variance = (sumSquaredEnergy / n_steps) - sumEnergy * sumEnergy / (n_steps * n_steps);
    return normalized ? variance / (latticeSize * latticeSize * latticeSize * latticeSize) : variance;
}

double IsingModel::calculateHeatCapacity(bool normalized) const {
    double meanEnergyVar = calculateEnergyVariance(true);
    double heatCapacity = meanEnergyVar / (kB * temp * temp);
    return normalized ? heatCapacity / (latticeSize * latticeSize) : heatCapacity;
}

double IsingModel::calculateSusceptibility(bool normalized) const {
    double meanMagVar = calculateMagnetizationVariance(true);
    double susceptibility = meanMagVar / (kB * temp);
    return normalized ? susceptibility / (latticeSize * latticeSize) : susceptibility;
}

void IsingModel::saveResults(const std::string& filename) const {
    std::ofstream outFile(filename);
    if (!outFile) {
        throw std::runtime_error("Could not open file for writing: " + filename);
    }

    outFile << "Lattice Size: " << latticeSize << "\n";
    outFile << "Temperature: " << temp << "\n";
    outFile << "Number of Steps: " << n_steps << "\n";
    outFile << "Average Energy: " << calculateEnergy() << "\n";
    outFile << "Average Magnetization: " << calculateMagnetization() << "\n";
    outFile << "Heat Capacity: " << calculateHeatCapacity() << "\n";
    outFile << "Susceptibility: " << calculateSusceptibility() << "\n";

    outFile.close();
}

void IsingModel::recordEvolution() {
    double en = currentEnergy();
    double mag = currentMagnetization();
    double En = calculateEnergy();
    double Mag = calculateMagnetization();
    evolutionData.push_back({static_cast<double>(n_steps), en, mag, En, Mag});
}

void IsingModel::saveEvolution(const std::string& filename) const {
    std::ofstream outFile(filename);
    if (!outFile) {
        throw std::runtime_error("Could not open file for writing: " + filename);
    }

    outFile << "Step,CurrentEnergy,CurrentMagnetization,AvgEnergy,AvgMagnetization\n";
    for (const auto& record : evolutionData) {
        outFile << record[0] << "," << record[1] << "," << record[2] << "," << record[3] << "," << record[4] << "\n";
    }

    outFile.close();
}