#include "classes.hpp"
#include "solvers.hpp"
#include "constants.hpp"


// Initialize base solver class
// dt in seconds
Solver::Solver(PenningTrap& trap_in, double dt_in) 
    : trap(trap_in), dt(dt_in * constants::second), steps(0) {}

// Simulate the system for a given number of steps
void Solver::simulate(int num_steps) {
    for (int i = 0; i < num_steps; i++) {
        step();
        if (record_positions_and_velocities) { // Allow disabling of recording for performance
            // Store positions of all particles after each step
            record_position();
            // Store velocities of all particles after each step
            record_velocity();
            if (!trap.get_modified_fields()) {
                record_energies();
            }
        }
        if (record_inside_trap_count) {
            record_trap_count();
        }
    steps++;
    }                  
}

void Solver::record_position() {
    arma::mat current_positions(3, trap.num_particles());
    for (int j = 0; j < trap.num_particles(); j++) {
        current_positions.col(j) = trap.get_particle(j).position;
    }
    positions.push_back(current_positions);
}
void Solver::record_velocity() {
    arma::mat current_velocities(3, trap.num_particles());
    for (int j = 0; j < trap.num_particles(); j++) {
        current_velocities.col(j) = trap.get_particle(j).velocity;
    }
    velocities.push_back(current_velocities);
}

void Solver::record_energies() {
    double U = trap.total_potential_energy();
    double K = trap.total_kinetic_energy();
    double E = trap.total_energy();
    potential_energies.push_back(U);
    kinetic_energies.push_back(K);
    total_energies.push_back(E);
}

void Solver::record_trap_count() {
    inside_trap_count.push_back(trap.num_inside_trap());
}

// Get recorded positions of particles
std::vector<arma::mat> Solver::get_positions() {
    return positions;
}

// Get recorded velocities of particles
std::vector<arma::mat> Solver::get_velocities() {
    return velocities;
}
std::vector<double> Solver::get_potential_energies() {
    return potential_energies;
}
std::vector<double> Solver::get_kinetic_energies() {
    return kinetic_energies;
}
std::vector<double> Solver::get_total_energies() {
    return total_energies;
}

std::vector<int> Solver::get_inside_trap_count() {
    return inside_trap_count;
}

// Get the current step count
int Solver::get_current_step() {
    return steps;
}

// Calculate total steps for a given total time
// total_time in seconds
int Solver::total_steps(double total_time) {
    return static_cast<int>(total_time * constants::second / dt);
}

void Solver::save(std::string base_filename) {
    arma::cube pos_cube = to_cube(positions);
    save_to_file(base_filename + "_positions.dat", pos_cube);
    arma::cube vel_cube = to_cube(velocities);
    save_to_file(base_filename + "_velocities.dat", vel_cube);
    if (!trap.get_modified_fields()) {
        arma::vec U = arma::vec(potential_energies);
        save_to_file(base_filename + "_potential_energies.dat", U);
        arma::vec K = arma::vec(kinetic_energies);
        save_to_file(base_filename + "_kinetic_energies.dat", K);
        arma::vec E = arma::vec(total_energies);
        save_to_file(base_filename + "_total_energies.dat", E);
    }
}

void Solver::save_counts(std::string base_filename) {
    std::string extension = ".dat";
    if (inside_trap_count.size() == 1) {
        // Reduced to last element if only one entry
        extension = ".red";
    }
    std::ofstream file(base_filename + "_inside_trap_count" + extension);
    for (size_t i = 0; i < inside_trap_count.size(); i++) {
        file << i << " " << inside_trap_count[i] << "\n";
    }
    file.close();
}

// Initialize Euler solver
EulerSolver::EulerSolver(PenningTrap& trap_in, double dt_in) 
    : Solver(trap_in, dt_in) {}

// Perform a single Euler integration step
void EulerSolver::step() {
    int n = trap.num_particles();
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec F = trap.total_force(i, steps * dt);
        arma::vec a = F / p.mass;
        p.position += p.velocity * dt;
        p.velocity += a * dt;
    }
}

// Initialize Runge-Kutta 4 solver
RK4Solver::RK4Solver(PenningTrap& trap_in, double dt_in) 
    : Solver(trap_in, dt_in) {}


// Perform a single RK4 integration step
void RK4Solver::step() {
    int n = trap.num_particles();
    std::vector<arma::vec> k1_v(n), k1_r(n);
    std::vector<arma::vec> k2_v(n), k2_r(n);
    std::vector<arma::vec> k3_v(n), k3_r(n);
    std::vector<arma::vec> k4_v(n), k4_r(n);

    std::vector<arma::vec> current_positions(n);
    std::vector<arma::vec> current_velocities(n);

    for (int i = 0; i < n; i++) {
        current_positions[i] = trap.get_particle(i).position;
        current_velocities[i] = trap.get_particle(i).velocity;
    }
    double t = steps * dt;
    // Calculate k1
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec F = trap.total_force(i, t);
        arma::vec a = F / p.mass;
        k1_v[i] = a * dt;
        k1_r[i] = p.velocity * dt;
        p.position += k1_r[i] / 2;
        p.velocity += k1_v[i] / 2;
    }
    // Calculate k2
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec F = trap.total_force(i, t + dt / 2);
        arma::vec a = F / p.mass;
        k2_v[i] = a * dt;
        k2_r[i] = p.velocity * dt;
        p.position = current_positions[i] + k2_r[i] / 2;
        p.velocity = current_velocities[i] + k2_v[i] / 2;
    }
    // Calculate k3
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec F = trap.total_force(i, t + dt / 2);
        arma::vec a = F / p.mass;
        k3_v[i] = a * dt;
        k3_r[i] = p.velocity * dt;
        p.position = current_positions[i] + k3_r[i];
        p.velocity = current_velocities[i] + k3_v[i];
    }
    // Calculate k4
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec F = trap.total_force(i, t + dt);
        arma::vec a = F / p.mass;
        k4_v[i] = a * dt;
        k4_r[i] = p.velocity * dt;
        p.position = current_positions[i];
        p.velocity = current_velocities[i];
    }
    // Update positions and velocities
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        p.position = current_positions[i] + (k1_r[i] + 2.0 * k2_r[i] + 2.0 * k3_r[i] + k4_r[i]) / 6.0;
        p.velocity = current_velocities[i] + (k1_v[i] + 2.0 * k2_v[i] + 2.0 * k3_v[i] + k4_v[i]) / 6.0;
    }
}

VelocityVerletSolver::VelocityVerletSolver(PenningTrap& trap_in, double dt_in) 
    : Solver(trap_in, dt_in) {
        // Initialize accelerations
        int n = trap.num_particles();
        accelerations.resize(n);
        for (int i = 0; i < n; i++) {
            Particle& p = trap.get_particle(i);
            arma::vec F = trap.total_force(i, 0.0);
            accelerations[i] = F / p.mass;
    }
}

// Perform a single Velocity Verlet integration step
void VelocityVerletSolver::step() {
    int n = trap.num_particles();
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        p.position += p.velocity * dt + 0.5 * accelerations[i] * dt * dt;
        arma::vec new_acceleration = trap.total_force(i, steps * dt) / p.mass;
        p.velocity += 0.5 * (accelerations[i] + new_acceleration) * dt;
        accelerations[i] = new_acceleration;
    }
}


// Initialize Analytical solver
AnalyticalSolver::AnalyticalSolver(PenningTrap& trap_in, double dt_in) 
    : Solver(trap_in, dt_in) {
    // Check for validity of analytical solution
    if (trap_in.get_coulomb_interaction()) {
        throw std::runtime_error("Analytical solution not valid with Coulomb interactions enabled.");
    }
    if (trap_in.get_modified_fields()) {
        throw std::runtime_error("Analytical solution not valid with modified external fields.");
    }
    for (int i = 0; i < trap_in.num_particles(); i++) {
        if (trap_in.get_particle(i).velocity(0) != 0 || trap_in.get_particle(i).velocity(2) != 0) {
            throw std::runtime_error("Analytical solution not valid for non-zero initial velocity in x or z direction.");
        }
        if (trap_in.get_particle(i).position(1) != 0) {
            throw std::runtime_error("Analytical solution not valid for non-zero initial position in y direction.");
        }
    }

    trap = trap_in;
    dt = dt_in * constants::second;
    steps = 0;
    int n = trap.num_particles();
    w_0.resize(n);
    w_z.resize(n);
    w_plus.resize(n);
    w_minus.resize(n);
    A_plus.resize(n);
    A_minus.resize(n);
    z_0.resize(n);

    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        double q = p.charge;
        double m = p.mass;
        double B0 = trap.B0;
        double V0 = trap.V0;
        double d = trap.d;

        w_0[i] = q * B0 / m;
        w_z[i] = std::sqrt(2 * q * V0 / (m * d * d));
        double discriminant = w_0[i] * w_0[i] - 2 * w_z[i] * w_z[i];
        if (discriminant < 0) {
            throw std::runtime_error("Analytical solution not valid: discriminant is negative.");
        }
        w_plus[i] = (w_0[i] + std::sqrt(discriminant)) / 2.0;
        w_minus[i] = (w_0[i] - std::sqrt(discriminant)) / 2.0;

        arma::vec r = p.position;
        arma::vec v = p.velocity;

        // Initial conditions
        double x0 = r(0);
        double y0 = r(1);
        double z0 = r(2);
        double vx0 = v(0);
        double vy0 = v(1);
        double vz0 = v(2);

        // Calculate amplitudes
        A_plus[i] = (vy0 + w_minus[i] * x0) / (w_plus[i] - w_minus[i]);
        A_minus[i] = -(vy0 + w_plus[i] * x0) / (w_plus[i] - w_minus[i]);
        z_0[i] = z0;
    }
}

BorisSolver::BorisSolver(PenningTrap& trap_in, double dt_in) 
    : Solver(trap_in, dt_in) {}

void BorisSolver::step() {
    int n = trap.num_particles();
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        arma::vec E = trap.external_E_field(p.position, steps * dt);
        E += trap.total_force_particles(i) / p.charge; // Add Coulomb force as effective E-field
        arma::vec B = trap.external_B_field(p.position, steps * dt);
        arma:: vec t = (p.charge * B / p.mass) * (dt / 2.0);
        arma::vec s = 2.0 * t / (1.0 + arma::dot(t, t));
        arma::vec v_minus = p.velocity + (p.charge * E / p.mass) * (dt / 2.0);
        arma::vec temp = arma::cross(v_minus, t); // Move to separate line to avoid compiler warning with -O3 flag due to inlining of cross
        arma::vec v_prime = v_minus + temp;
        temp = arma::cross(v_prime, s);
        arma::vec v_plus = v_minus + temp;
        p.velocity = v_plus + (p.charge * E / p.mass) * (dt / 2.0);
        p.position += p.velocity * dt;
    }
}

// Perform a single Analytical "integration" step
void AnalyticalSolver::step() {
    int n = trap.num_particles();
    double t = steps * dt;
    for (int i = 0; i < n; i++) {
        Particle& p = trap.get_particle(i);
        double x = -A_plus[i] * std::cos(w_plus[i] * t) - A_minus[i] * std::cos(w_minus[i] * t);
        double y = A_plus[i] * std::sin(w_plus[i] * t) + A_minus[i] * std::sin(w_minus[i] * t);
        double z = z_0[i] * std::cos(w_z[i] * t);

        double vx = +A_plus[i] * w_plus[i] * std::sin(w_plus[i] * t) + A_minus[i] * w_minus[i] * std::sin(w_minus[i] * t);
        double vy = A_plus[i] * w_plus[i] * std::cos(w_plus[i] * t) + A_minus[i] * w_minus[i] * std::cos(w_minus[i] * t);
        double vz = -z_0[i] * w_z[i] * std::sin(w_z[i] * t);

        p.position = arma::vec({x, y, z});
        p.velocity = arma::vec({vx, vy, vz});
    }
}


// Convert vector of matrices to a 3D cube for easier data handling
arma::cube to_cube(const std::vector<arma::mat>& vec) {
    if (vec.empty()) {
        return arma::cube();
    }
    int rows = vec[0].n_rows;
    int cols = vec[0].n_cols;
    int slices = vec.size();
    arma::cube cube_data(rows, cols, slices);
    for (size_t i = 0; i < vec.size(); ++i) {
        cube_data.slice(i) = vec[i];
    }
    return cube_data;
}

// Save cube data to a file
void save_to_file(const std::string& filename, const arma::cube& data) {
    data.save(filename, arma::raw_ascii);
}

void save_to_file(const std::string& filename, const arma::vec& data) {
    arma::mat mat_data = arma::conv_to<arma::mat>::from(data);
    mat_data.save(filename, arma::raw_ascii);
}