#include <iostream>
#include <iomanip>
#include <string>
#include <sstream>
#include <fstream>
#include <functional>
#include <vector>
#include <limits>
#include <utility>
#include "eigen/Dense"
#include "eigen/Eigen"

//Class Representing Shallow-Water Problem
class Shallow_Water_Solver {
public:

    //Default Constructors and Assignment
    Shallow_Water_Solver() = default;
    Shallow_Water_Solver(const Shallow_Water_Solver&) = default;
    Shallow_Water_Solver(Shallow_Water_Solver&&) = default;
    Shallow_Water_Solver& operator=(const Shallow_Water_Solver&) = default;
    Shallow_Water_Solver& operator=(Shallow_Water_Solver&&) = default;

    //Constructors with Real Information
    Shallow_Water_Solver(double x_min_in, double x_max_in, int nx_in,
                         double y_min_in, double y_max_in, int ny_in,
                         const std::function<Eigen::Vector3d(double,double)>& initial_condition) :
            SolutionVector(3*nx_in*ny_in),
            ResidualVector(3*nx_in*ny_in),
            x_min(x_min_in),
            x_max(x_max_in),
            nx(nx_in),
            delta_x((x_max-x_min)/static_cast<double>(nx+1)),
            y_min(y_min_in),
            y_max(y_max_in),
            ny(ny_in),
            delta_y((y_max-y_min)/static_cast<double>(ny+1)),
            current_time(0.0)
    {
        if(x_min > x_max || nx <= 0 || y_min > y_max || ny <= 0) {
            throw std::runtime_error("Invalid inputs to solver.");
        }
        set_initial_conditions(initial_condition);
    }

    auto num_x() const {return nx;} //Number of Nodes in the X Direction
    auto num_y() const {return ny;} //Number of Nodes in the Y Direction
    auto dx() const {return delta_x;} //Delta X
    auto dy() const {return delta_y;} //Delta Y

    //Position of i-j Node
    Eigen::Vector2d node(int i, int j) const {
        while(i<0) i += nx;
        auto index_i = i%nx;
        while(j<0) j += ny;
        auto index_j = j%ny;
        return {x_min + static_cast<double>(index_i)*delta_x + (i/nx)*(x_max-x_min),
                y_min + static_cast<double>(index_j)*delta_y + (j/ny)*(y_max-y_min)};
    }

    auto& U() {return SolutionVector;} //Full Solution Vector
    const auto& U() const {return SolutionVector;} //Full Solution Vector

    //Solution at i-j Node. Returns Eigen Segment. Desired SubVector Reference
    auto U(int i, int j) {
        auto index = global_solution_index(i,j);
        return SolutionVector.segment<3>(index);
    }

    //Solution at i-j Node. Returns Eigen Segment. Desired SubVector Reference
    auto U(int i, int j) const {
        auto index = global_solution_index(i,j);
        return SolutionVector.segment<3>(index);
    }

    auto& dUdt() {return ResidualVector;} //Full Residual Vector
    const auto& dUdt() const {return ResidualVector;} //Full Residual Vector

    //Solution at i-j Node. Returns Eigen Segment. Desired SubVector Reference
    auto dUdt(int i, int j) {
        auto index = global_solution_index(i,j);
        return ResidualVector.segment<3>(index);
    }

    //Solution at i-j Node. Returns Eigen Segment. Desired SubVector Reference
    auto dUdt(int i, int j) const {
        auto index = global_solution_index(i,j);
        return ResidualVector.segment<3>(index);
    }

    auto time() const {return current_time;} //Current Time

    //Force X
    Eigen::Vector3d Fx(const Eigen::Vector3d& U) {
        return {U[1], U[1]*U[1]/U[0] + 0.5*g*U[0]*U[0], U[1]*U[2]/U[0]};
    }

    //Force Y
    Eigen::Vector3d Fy(const Eigen::Vector3d& U) {
        return {U[2], U[2]*U[1]/U[0], U[2]*U[2]/U[0] + 0.5*g*U[0]*U[0]};
    }

    //Max Wave Speeds in X and Y Directions.
    Eigen::Vector2d max_wavespeeds(const Eigen::Vector3d& U) {
        double a  = sqrt(U[0]*g);
        double ux = U[1]/U[0];
        double uy = U[2]/U[0];
        return {fabs(ux)+a, fabs(uy)+a};
    }

    void time_march(double final_time, double CFL); //Time March to Time

private:

    // Index in Solution Vector. Storing Location for Node i-j Entries
    int global_solution_index(int i, int j) const {
        while(i<0) i += nx;
        auto index_i = i%nx;
        while(j<0) j += ny;
        auto index_j = j%ny;
        return 3 * (index_i + nx*index_j);
    }

    //Set Initial Conditions
    void set_initial_conditions(const std::function<Eigen::Vector3d(double, double)>& initial_condition) {

        for(int i = 0; i < nx; ++i) {
            for(int j = 0; j < ny; ++j) {
                auto n = node(i,j);
                U(i,j) = initial_condition(n.x(), n.y());
            }
        }

    }

    Eigen::VectorXd SolutionVector; //Solution vector (U)
    Eigen::VectorXd ResidualVector; // Residual Vector (dUdt)
    const double x_min; //Minimum Value of X
    const double x_max; //Maximum value of X
    const int nx; //Number of Nodes in the X Direction
    const double delta_x; //Spacing of Nodes in the X Direction
    const double y_min; //Minimum Value of Y
    const double y_max; //Maximum Value of Y
    const int ny; //Number of Nodes in the Y Direction
    const double delta_y; //Spacing of Nodes in the Y Direction
    double current_time; //Solution Time
    constexpr static double g = 9.81; //Gravity

};

//Time March to Time. Local Lax-Friedrichs
void Shallow_Water_Solver::time_march(double final_time, double CFL) {
    std::cout << "Time marching from t = " << current_time
              << " to t = " << final_time
              << ".  Total difference = " << final_time-current_time << ".\n";
    const double time_tolerance = 1.0e-12;
    const double min_length = std::min(dx(), dy());
    while(current_time < final_time - time_tolerance) {
        auto dt = std::numeric_limits<double>::max();
        dUdt().fill(0.0);
        for(int i = 0; i < num_x(); ++i) {
            for(int j = 0; j < num_y(); ++j) {
                Eigen::Vector3d Um = U(i  , j  );
                Eigen::Vector3d Ul = U(i-1, j  );
                Eigen::Vector3d Ur = U(i+1, j  );
                Eigen::Vector3d Ub = U(i  , j-1);
                Eigen::Vector3d Ut = U(i  , j+1);
                Eigen::Vector2d lambda = max_wavespeeds(Um);
                auto max_lambda = std::max(lambda[0], lambda[1]);
                dt = std::min(dt, CFL*min_length/max_lambda);
                dUdt(i,j) = (Fx(Ul) - Fx(Ur) + lambda[0]*(Ul - 2.0*Um + Ur)) / (2.0*delta_x)
                            +(Fy(Ub) - Fy(Ut) + lambda[1]*(Ub - 2.0*Um + Ut)) / (2.0*delta_y);
            }
        }

        if(current_time + dt > final_time) {
            dt = final_time - current_time;
        }

        U()  += dt*dUdt();
        current_time += dt;
        std::cout << "Time = " << std::setw(10) << current_time << '\n';
    }

    std::cout << "Time marching done.\n";
}

//Write to VTK
void write_to_VTK(const Shallow_Water_Solver& solver,
                  const std::string& filename) {
    std::ofstream fout(filename);
    if(!fout) {
        throw std::runtime_error("Could not open file: " + filename);
    }

    std::cout << "Writing output to file: " << filename << '\n';
    const auto num_nodes = (solver.num_x()+1)*(solver.num_y()+1);
    fout << "# vtk DataFile Version 2.0\n"
         << "Shallow-Water equations solution\n"
         << "ASCII\n"
         << "DATASET STRUCTURED_GRID\n"
         << "DIMENSIONS " << solver.num_x()+1 << " " << solver.num_y()+1 << " 1\n"
         << "POINTS " << num_nodes << " double\n";
    for(int j = 0; j <= solver.num_y(); ++j) {
        for(int i = 0; i <= solver.num_x(); ++i) {
            auto n = solver.node(i,j);
            auto U = solver.U(i,j);
            fout << n.x() << " " << n.y() << " " << U[0] << '\n';
        }
    }

    fout << "\nPOINT_DATA " << num_nodes
         << "\nSCALARS h double 1\nLOOkUP_TABLE default\n";
    for(int j = 0; j <= solver.num_y(); ++j) {
        for(int i = 0; i <= solver.num_x(); ++i) {
            fout << solver.U(i,j)[0] << '\n';
        }
    }

    fout << "\nVECTORS u double\n";
    for(int j = 0; j <= solver.num_y(); ++j) {
        for(int i = 0; i <= solver.num_x(); ++i) {
            auto ux = solver.U(i,j)[1]/solver.U(i,j)[0];
            auto uy = solver.U(i,j)[2]/solver.U(i,j)[0];
            fout << ux << " " << uy << " 0.0\n";
        }
    }
}

//Make Movie
void make_movie(Shallow_Water_Solver& solver,
                const double final_time,
                const double CFL,
                const int number_of_frames,
                const std::string& filename_base) {
    if(number_of_frames < 2) {
        throw std::runtime_error("make_movie requires at least two frames.");
    }

    const auto frame_time = (final_time-solver.time())/static_cast<double>(number_of_frames-1);
    const auto build_filename = [](const std::string& filename_base, int index) {
        std::stringstream filename_ss;
        filename_ss <<  filename_base << "_" << std::setfill('0') << std::setw(5) << index << ".vtk";
        return filename_ss.str();
    };

    write_to_VTK(solver, build_filename(filename_base, 0));
    for(int i = 1; i<number_of_frames; ++i) {
        double target_time = static_cast<double>(i)*frame_time;
        solver.time_march(target_time, CFL);
        write_to_VTK(solver, build_filename(filename_base, i));
    }
}

//Main
int main() {
    std::cout << "---------------------------------------\n"
              << " |              Lab 6                |\n"
              << " |           Shallow-water           |\n"
              << "---------------------------------------\n";
    double x_min = -1.0;
    double x_max =  1.0;
    int       nx =  500;
    double y_min = -1.0;
    double y_max =  1.0;
    int       ny =  500;
    auto initial_condition = [] (double x, double y) {
        Eigen::Vector3d U;
        U[0] = 1.0;
        U[1] = 0.0;
        U[2] = 0.0;
        if(fabs(x)<0.2 && fabs(y)<0.2) {
            U *= 1.2;
        }
        return U;
    };

    auto solver = Shallow_Water_Solver(x_min, x_max, nx, y_min, y_max, ny, initial_condition);
    make_movie(solver, 2.5, 0.5, 750, "movie");
    return 0;
}