/* Distributed under the Apache License, Version 2.0.
   See accompanying NOTICE file for details.*/

#pragma once

#include "LBMExports.h"
#include "LogUtils.h"

extern "C"
LBM_DECL bool cuda_device_check(Logger & log, int use_device = -1);

namespace LBM
{
  class Run;

  enum class BoundaryCondition { Pressure = 0, Flow = 1 };

  enum class GeometryType { Custom = 0, Brick = 1, Cylinder = 2 };

  // Old school enum as we are indexing arrays with this value
  enum ModelType { Preconditioner = 0, Fluid = 1, Thermal = 2 };

  enum class BoundaryTypes : int
  {
    Inactive = -1,
    Wall = 0,
    Inlet = 100,
    Outlet = 200
  };

#ifdef _MSC_VER
#pragma warning( push )
#pragma warning( disable : 4251 )
  // Explicit template exports.
  LBM_EXP template class LBM_DECL std::allocator<LBM::BoundaryTypes>;
  LBM_EXP template class LBM_DECL std::vector<LBM::BoundaryTypes, std::allocator<LBM::BoundaryTypes>>;
  LBM_EXP template class LBM_DECL std::allocator<std::pair<const int, LBM::BoundaryTypes>>;
  LBM_EXP template class LBM_DECL std::map<int, LBM::BoundaryTypes, std::less<int>, std::allocator<std::pair<const int, LBM::BoundaryTypes>>>;
#pragma warning( pop )
#endif

  struct LBM_DECL Configuration
  {
    std::string name;
    // You can provide any values of any unit
    // As long as they are all consistent
    // ie.e All lengths are in m, time in s
    //   - density in kg_per_m3
    //   - viscosity in m2_per_s
    //   - speed_of_sound in m_per_s
    //   - grid_spacing in m

    float atmospheric_pressure=0;
    float density=0;
    float fluid_temperature=0;
    float initial_velocity=0;
    float mach_number = 0;
    float relaxation_time = 0;
    float relaxation_time_temperature = 0;
    float reynolds_number = 0;
    float speed_of_sound=0;
    float temporal_step = 0;
    float thermal_coefficient=0;
    float viscosity=0;
    float unitless_viscosity=0;
    float target_flow_rate=0;
    float wall_temperature=0;
    // Boundary Conditions
    BoundaryCondition inlet_boundary_condition;
    float inlet_pressure=0;
    float inlet_volumetric_flow_rate=0;
    BoundaryCondition outlet_boundary_condition;
    float outlet_pressure=0;
    float outlet_volumetric_flow_rate=0;
    // Solver Execution
    size_t max_iterations;
    size_t print_stride;
    size_t print_iteration;
    float tolerance;

    // General Run Flags
    bool expanded_results = true;
    bool setup_only = false;
    bool write_all_timesteps = false;

    void ToString(std::ostream& str) const;
  };

  struct LBM_DECL Input
  {
    Input();
    ~Input() = default;

    float        origin[3];
    size_t       dimensions[3];
    ModelType    model_type;

    // Geometry Specification
    // If a non-custom geometry type is specified
    // A lattice will be created based on the type parameters
    GeometryType geometry_type;
    float        brick_x_length;
    float        brick_y_length;
    float        brick_z_length;
    float        cylinder_radius;
    float        cylinder_length;
    float        grid_spacing;

    // Dimensionless values for computation
    // These will get set in LBM based on configuration data
    float        bc_accuracy = 0;
    float        density = 0;
    float        essential_condition = 0;
    float        gamma = 0;
    float        knudsen_number = 0;
    float        mach_number = 0;
    float        reynolds_max = 0;
    float        relaxation_time = 0;
    float        relaxation_time_temperature = 0;
    float        reynolds_grid_number = 0;
    float        reynolds_number = 0;
    float        spatial_step = 0;
    float        speed_of_sound = 0;
    float        system_length = 0;
    float        system_velocity = 0;
    float        temporal_step = 0;
    float        viscosity = 0;
    float        thermal_coefficient = 0;

    bool         return_distribution_functions;
    bool         initialize_distribution_functions_from_cont_mech;
    bool         pre_conditioning;

    // You have two options for providing labels to LBM
    // 1. A vector of LBM labels
    std::vector<BoundaryTypes> labels;
    // 2. Your own vector of labels, with a mapping to LBM
    //   - This is provided to store your labels with LBM
    //   - Source labels can be modified by some LBM preprocessors
    std::vector<int> source_labels;
    std::map<int, BoundaryTypes> source_to_lbm_label_map;
  };

  class LBM_DECL Output : protected Loggable
  {
  public:
    Output(Run const& run);
    ~Output();

    bool Allocate();
    void Deallocate();
    bool CountCellTypes();
    size_t Index(size_t x, size_t y, size_t z) const;

    size_t GetNumCells() const { return num_cells; }
    size_t GetNumWallCells() const { return num_wall_cells; }
    size_t GetNumInflowCells() const { return num_inflow_cells; }
    size_t GetNumOutflowCells() const { return num_outflow_cells; }
    size_t GetNumInactiveCells() const { return num_inactive_cells; }
    size_t GetNumInteriorCells() const { return num_interior_cells; }
    size_t GetNumUnknownCells() const { return num_unknown_cells; }

    //    float GetReferenceVelocity() const { return characteristic_velocity; }
    float GetInletVelocity() const { return inlet_velocity; }
    float GetOutletVelocity() const { return outlet_velocity; }
    float GetInletArea() const { return inlet_area; }
    float GetOutletArea() const { return outlet_area; }
    float GetOutletDiameter() const { return average_diameter; }
    float GetAverageArea() const { return average_area; }
    float GetMassFlow() const { return mass_flow; }
    float GetPressureDrop() const { return pressure_drop; }
    float GetVolumetricFlowRate() const { return volumetric_flow_rate; }

    const int* GetLabels() const { return labels; }
    int  GetLabel(size_t x, size_t y, size_t z) const { return labels[Index(x, y, z)]; }
    const float* GetPressure() const { return pressure_flows_and_stress; }
    float  GetPressure(size_t x, size_t y, size_t z) const { return pressure_flows_and_stress[Index(x, y, z)]; }
    const float* GetXFlow() const { return &pressure_flows_and_stress[num_cells]; }
    float  GetXFlow(size_t x, size_t y, size_t z) const { return pressure_flows_and_stress[num_cells + Index(x, y, z)]; }
    const float* GetYFlow() const { return &pressure_flows_and_stress[num_cells * 2]; }
    float  GetYFlow(size_t x, size_t y, size_t z) const { return pressure_flows_and_stress[num_cells * 2 + Index(x, y, z)]; }
    const float* GetZFlow() const { return &pressure_flows_and_stress[num_cells * 3]; }
    float  GetZFlow(size_t x, size_t y, size_t z) const { return pressure_flows_and_stress[num_cells * 3 + Index(x, y, z)]; }
    const float* GetTemperature() const { return temperature; }
    float  GetTemperature(size_t x, size_t y, size_t z) const { return temperature[Index(x, y, z)]; }
    const float* GetWallShearStress() const { return &wall_shear_stress[num_cells]; }
    float  GetWallShearStress(size_t x, size_t y, size_t z) const { return wall_shear_stress[Index(x, y, z)]; }
    // Expanded results, may not be available
    const float* GetLegacyStress() const { return wall_shear_stress; }
    const float* GetXYStress() const { return &pressure_flows_and_stress[num_cells * 4]; }
    const float* GetXZStress() const { return &pressure_flows_and_stress[num_cells * 5]; }
    const float* GetYZStress() const { return &pressure_flows_and_stress[num_cells * 6]; }
    const float* GetXXStress() const { return &pressure_flows_and_stress[num_cells * 7]; }
    const float* GetYYStress() const { return &pressure_flows_and_stress[num_cells * 8]; }
    const float* GetZZStress() const { return &pressure_flows_and_stress[num_cells * 9]; }

    // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    // !! Raw data values, handle with caution !!
    // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    // Specify nr. of continuum components (p,ux,uy,uz) and viscous stress components (Sigma_xy, Sigma_xz, Sigma_yz)
    const size_t nQ = 4;
    const size_t nS = 6;
    // Specify nr. of spatial labels (exterior, interior, wall, inlet and outlet)
    const size_t nBound = 5;

    size_t num_cells = 0;
    // Count up label types
    size_t num_wall_cells = 0;
    size_t num_inflow_cells = 0;
    size_t num_outflow_cells = 0;
    size_t num_inactive_cells = 0;
    size_t num_interior_cells = 0;
    size_t num_unknown_cells = 0;

    size_t num_preconditioner_iterations = 0;
    size_t num_fluid_iterations = 0;
    size_t num_thermal_iterations = 0;

    // Secondary parameters to be computed
    float  inlet_velocity = 0;
    float  inlet_pressure = 0;
    float  outlet_pressure = 0;
    float  inlet_area = 0;
    float  outlet_velocity = 0;
    float  outlet_area = 0;
    float  average_area = 0;
    float  average_diameter = 0;
    float  axial_length = 0;
    float  mass_flow = 0;
    float  pressure_drop = 0;
    float  volumetric_flow_rate = 0;

    // Array Lengths
    size_t Q_num_cell_values = 0;   // Per q array passed into microLBM
    size_t T_num_cell_values = 0;   // Per q array passed into microLBM
    size_t num_boundary_values = 0;
    size_t nPDF = 0;

    int* labels = nullptr;
    float* pressure_flows_and_stress = nullptr;
    float* wall_shear_stress = nullptr;
    float* temperature = nullptr;
    float* fPDF = nullptr;//host memory for PDFs
    float* fnPDF = nullptr;//host memory for non-equilibrium PDFs
    float* boundary_conditions = nullptr;
    float* previous_results = nullptr; // Used to compute convergance error

    std::vector<float> pre_conditioning_iteration_error;
    std::vector<float> pre_conditioning_iteration_timing_ms;
    std::vector<float> fluid_iteration_error;
    std::vector<float> fluid_iteration_timing_ms;
    std::vector<float> temperature_iteration_error;
    std::vector<float> temperature_iteration_timing_ms;

    float compute_time_s = 0;

  protected:
    Run const& run;
  };

  class LBM_DECL Run : public Loggable
  {
  public:
    Run(Logger* logger = nullptr);
    Run(std::string const& logfile);
    virtual ~Run();

    bool Compute();

    Input                 in;
    Output                out;
    Configuration         cfg;
  protected:
    bool GetNormals();
  };

  struct LBM_DECL FlowMetrics
  {
    float flow_rate;
    float resistance;
    float heat_flux;
    float wall_shear;

    void ToString(std::ostream& str) const;
  };

  // LBM direction encoding. 3 letter variable formed from O=zero, P=plus, M=minus
  // encode in 6 bits zzyyxx with less significant bit associated with P direction
  // LDM model indices are used as bit masks in the cuda code (LBM_Stream)
  #define OOO  0  // 00 00 00                 0
  #define OOP  1  // 00 00 01                 1
  #define OOM  2  // 00 00 10                 2
  #define OPO  4  // 00 01 00                 3
  #define OPP  5  // 00 01 01                 4
  #define OPM  6  // 00 01 10                 5
  #define OMO  8  // 00 10 00                 6
  #define OMP  9  // 00 10 01                 7
  #define OMM 10  // 00 10 10                 8
  #define POO 16  // 01 00 00                 9
  #define POP 17  // 01 00 01                10
  #define POM 18  // 01 00 10                11
  #define PPO 20  // 01 01 00                12
  #define PMO 24  // 01 10 00                13
  #define MOO 32  // 10 00 00                14
  #define MOP 33  // 10 00 01                15
  #define MOM 34  // 10 00 10                16
  #define MPO 36  // 10 01 00                17
  #define MMO 40  // 10 10 00                18
  // Direction codes in order of increasing distance from current node
  //                Direction    0       1       2       4       8       16      32        5      10         6        9       20      40
  //                           zyx     zyx     zyx      zyx     zyx      zyx     zyx      zyx     zyx       zyx      zyx      zyx     zyx
  static int VelocityDirectionMasks[19] = { OOO,OOP,OOM,OPO,OPP,OPM,OMO,OMP,OMM,POO,POP,POM,PPO,PMO,MOO,MOP,MOM,MPO,MMO };
  static int VelocityDirections[19][3] = { {0,0,0},{0,0,1},{0,0,-1},{0,1,0},{0,-1,0},{1,0,0},{-1,0,0},{0,1,1},{0,-1,-1},{0,1,-1},{0,-1,1},{1,1,0},{-1,-1,0},{1,-1,0},{-1,1,0},{1,0,1},{-1,0,-1},{1,0,-1},{-1,0,1} };
  static int StreamDirections[19] = { 0,  1,  1,  1,  2,  2,  1,  2,  2,  1,  2,  2,  2,  2,  1,  2,  2,  2,  2 };

  constexpr int nPDFmodels = 19;
  constexpr int nContmodels[3] = { 4, 10, 4 };
}

inline std::ostream& operator<< (std::ostream& out, LBM::Configuration const& cfg)
{
  cfg.ToString(out);
  return out;
}
inline std::ostream& operator<< (std::ostream& out, LBM::FlowMetrics const& fm)
{
  fm.ToString(out);
  return out;
}
