import enum
import json
import typing
import dataclasses


# === DEFINE ALL ENUMERATIONS === #


class PatchGeometry(enum.Enum):
    D1_LINE_SEGMENT =  1
    D2_CIRCLE       =  2
    D2_RECTANGLE    =  3
    D2_SWEEP_LINE   =  4
    D2_ELLIPSE      =  5
    D2_VORTEX       =  6
    D2_ANALYTICAL   =  7
    D3_SPHERE       =  8
    D3_CUBOID       =  9
    D3_CYLINDER     = 10
    D3_SWEEP_PLANE  = 11
    D3_ELLIPSOID    = 12
    D3_ANALYTICAL   = 13


class FluxLimiter(enum.Enum):
    MINMOD     = 1
    MC         = 2
    OSPRE      = 3
    SUPERBEE   = 4
    SWEBY      = 5
    VAN_ALBADA = 6
    VAN_LEER   = 7


class MulticomponentModel(enum.Enum):
    GAMMA_PI_INF = 1
    EQUATION_5   = 2
    equation_6   = 3


class BubbleModel(enum.Enum):
    GILMORE          = 1
    KELLER_MIKSIS    = 2
    RAYLEIGH_PLESSET = 3


class ThermalModel(enum.Enum):
    ADIABATIC  = 1
    ISOTHERMAL = 2
    TRANSFER   = 3


class BoundaryCondition(enum.Enum):
    PERIODIC                           =  -1
    REFLECTIVE                         =  -2
    GHOST_CELL_EXTRAPOLATION           =  -3
    RIEMANN_EXTRAPOLATION              =  -4
    SLIP_WALL                          =  -5
    NON_REFLECTING_SUBSONIC_BUFFER     =  -6
    NON_REFLECTING_SUBSONIC_INFLOW     =  -7
    NON_REFLECTING_SUBSONIC_OUTFLOW    =  -8
    FORCE_FREE_SUBSONIC_OUTFLOW        =  -9
    CONSTANT_PRESSURE_SUBSONIC_OUTFLOW = -10
    SUPERSONIC_INFLOW                  = -11
    SUPERSONIC_OUTFLOW                 = -12


class DatabaseFormat(enum.Enum):
    SILO_HDF5 = 1
    BINARY    = 2


class FloatingPrecision(enum.Enum):
    SINGLE = 1
    DOUBLE = 2


class RiemannSolver(enum.Enum):
    HLL   = 1
    HLLC  = 2
    EXACT = 3


class WaveSpeedEstimation(enum.Enum):
    DIRECT            = 1
    PRESSURE_VELOCITY = 2


class AverageStateEvaluation(enum.Enum):
    ROE_MEAN        = 1
    ARITHMETIC_MEAN = 2


class TimeStepper(enum.Enum):
    RUNGE_KUTTA_1 = 1
    RUNGE_KUTTA_2 = 2
    RUNGE_KUTTA_3 = 3
    RUNGE_KUTTA_4 = 4
    RUNGE_KUTTA_5 = 5


class WenoVariables(enum.Enum):
    CONSERVATIVE = 1
    PRIMITIVE    = 2


# === DEFINE UTILITY CLASSES === #


@dataclasses.dataclass
class Point:
    x: float = dataclasses.field(default=0)
    y: float = dataclasses.field(default=0)
    z: float = dataclasses.field(default=0)


class Vector(Point):
    pass


# ===    DEFINE CLASSES    === #
# === COMPUTATIONAL DOMAIN === #


@dataclasses.dataclass
class AxisDomain:
    begin: float = dataclasses.field(default=0)
    end:   float = dataclasses.field(default=0)


@dataclasses.dataclass
class AxisStretch(AxisDomain):
    stretch: bool  = dataclasses.field(default=False)
    rate:    float = dataclasses.field(default=0)


@dataclasses.dataclass
class Stretch:
    x: AxisStretch = dataclasses.field(default=AxisStretch())
    y: AxisStretch = dataclasses.field(default=AxisStretch())
    z: AxisStretch = dataclasses.field(default=AxisStretch())


@dataclasses.dataclass
class Cells:
    x: int
    y: int = dataclasses.field(default=0)
    z: int = dataclasses.field(default=0)


@dataclasses.dataclass
class SpacialDomain:
    x: AxisDomain
    y: AxisDomain = dataclasses.field(default=AxisDomain(begin=0, end=0))
    z: AxisDomain = dataclasses.field(default=AxisDomain(begin=0, end=0))


@dataclasses.dataclass
class Time:
    end:   int
    save:  int
    dt:    float
    start: int   = dataclasses.field(default=0)


@dataclasses.dataclass
class ComputationalDomain:
    cells:        Cells
    domain:       SpacialDomain
    time:         Time
    cyl_coord:    bool    = dataclasses.field(default=False)
    stretch:      Stretch = dataclasses.field(default=Stretch())


# ===        FLUIDS        === #


@dataclasses.dataclass
class Fluid:
    gamma:   float      = None
    pi_inf:  float      = None
    mul0:    typing.Any = None
    ss:      typing.Any = None
    gamma_v: typing.Any = None
    M_v:     typing.Any = None
    mu_v:    typing.Any = None
    k_v:     typing.Any = None
    G:       typing.Any = None


# === SIMULATION ALGORITHM === #


@dataclasses.dataclass
class WenoParameters:
    order:     int
    variables: WenoVariables
    mapped:    bool
    mp:        bool
    epsilon:   float = dataclasses.field(default=1e-16)


@dataclasses.dataclass
class AxisBoundaryCondition:
    begin: BoundaryCondition
    end:   BoundaryCondition


@dataclasses.dataclass
class BoundaryConditions:
    x: AxisBoundaryCondition
    y: AxisBoundaryCondition = dataclasses.field(default=None)
    z: AxisBoundaryCondition = dataclasses.field(default=None)


@dataclasses.dataclass
class SimulationAlgorithm:
    model:          MulticomponentModel
    weno:           WenoParameters
    boundary:       BoundaryConditions
    riemann_solver: RiemannSolver
    time_stepper:   TimeStepper
    wave_speeds:    WaveSpeedEstimation
    avg_state:      AverageStateEvaluation
    mixture_err:    bool = dataclasses.field(default=False)
    mpp_lim:        bool = dataclasses.field(default=False)
    adv_alphan:     bool = dataclasses.field(default=False) 
    hypoelasticity: bool = dataclasses.field(default=False)
    null_weights:   bool = dataclasses.field(default=False)


# ===   PATCH PARAMETERS   === #


@dataclasses.dataclass
class Patch:
    geometry:        PatchGeometry
    radius:          float         = None
    radii:           typing.Any    = None
    epsilon:         typing.Any    = None
    beta:            typing.Any    = None
    smooth_patch_id: int           = None
    alpha:           typing.Any    = None
    alpha_rho:       typing.Any    = None
    smooth_coeff:    float         = None
    rho:             typing.Any    = None
    pressure:        float         = None
    gamma:           float         = None
    pi_inf:          float         = None
    r0:              typing.Any    = None
    v0:              typing.Any    = None
    p0:              typing.Any    = None
    m0:              typing.Any    = None
    centroid:        Point         = None
    length:          Vector        = None
    radii:           Vector        = None
    normal:          Vector        = None
    smoothen:        bool          = dataclasses.field(default=False)
    velocity:        Vector        = dataclasses.field(default=Vector(x=0,y=0,z=0))

    # TODO: Add for arho_id in range(1, 10+1):

    alter_patch:     typing.Any = dataclasses.field(default=False)


# ===  DATABASE STRUCTURE  === #


@dataclasses.dataclass
class DatabaseWrite:
    probe:      bool = dataclasses.field(default=False)
    cons_vars:  bool = dataclasses.field(default=False)
    prim_vars:  bool = dataclasses.field(default=False)
    rho:        bool = dataclasses.field(default=False)
    E:          bool = dataclasses.field(default=False)
    pres:       bool = dataclasses.field(default=False)
    gamma:      bool = dataclasses.field(default=False)
    heat_ratio: bool = dataclasses.field(default=False)
    pi_inf:     bool = dataclasses.field(default=False)
    pres_inf:   bool = dataclasses.field(default=False)
    c:          bool = dataclasses.field(default=False)
    schlieren:  bool = dataclasses.field(default=False)
    #TODO: Some (i) params


@dataclasses.dataclass
class DatabseStructure:
    precision:      FloatingPrecision 
    fd_order:       int
    write:          DatabaseWrite     = dataclasses.field(default=DatabaseWrite())
    alt_soundspeed: bool              = dataclasses.field(default=False)
    parallel_io:    bool              = dataclasses.field(default=False)
    format:         DatabaseFormat    = dataclasses.field(default=DatabaseFormat.SILO_HDF5)


# ===       BUBBLES       === #


@dataclasses.dataclass
class Bubbles:
    bubbles:    bool = None
    model:      BubbleModel = None
    polytropic: bool = None
    thermal:    ThermalModel = None
    R0ref:      float = None
    number:     int   = None
    cavitation: float = None
    weber:      float = None
    Re_inv:     float = None
    mu_10:      float = None
    ss:         float = None
    pv:         float = None
    gamma_v:    float = None
    M_v:        float = None
    mu_v:       float = None
    k_v:        float = None
    qbmm:       bool  = None


# === DEFINE A CASE === #


@dataclasses.dataclass
class Case:
    domain:    ComputationalDomain
    fluids:    typing.List[Fluid]
    patches:   typing.List[Patch]
    bubbles:   Bubbles
    database:  DatabseStructure
    algorithm: SimulationAlgorithm

    def __repr__(self) -> str:
        return json.dumps(self.to_json(), indent=4)

    def to_json(self) -> dict:
        def handle_enums(data):
            def convert_value(obj):
                if isinstance(obj, enum.Enum):
                    return obj.value

                return obj

            return dict((key, convert_value(value)) for key, value in data)

        return dataclasses.asdict(self, dict_factory=handle_enums)


import enum
import json
import typing
import dataclasses


# === DEFINE ALL ENUMERATIONS === #


class PatchGeometry(enum.Enum):
    D1_LINE_SEGMENT =  1
    D2_CIRCLE       =  2
    D2_RECTANGLE    =  3
    D2_SWEEP_LINE   =  4
    D2_ELLIPSE      =  5
    D2_VORTEX       =  6
    D2_ANALYTICAL   =  7
    D3_SPHERE       =  8
    D3_CUBOID       =  9
    D3_CYLINDER     = 10
    D3_SWEEP_PLANE  = 11
    D3_ELLIPSOID    = 12
    D3_ANALYTICAL   = 13


class FluxLimiter(enum.Enum):
    MINMOD     = 1
    MC         = 2
    OSPRE      = 3
    SUPERBEE   = 4
    SWEBY      = 5
    VAN_ALBADA = 6
    VAN_LEER   = 7


class MulticomponentModel(enum.Enum):
    GAMMA_PI_INF = 1
    EQUATION_5   = 2
    equation_6   = 3


class BubbleModel(enum.Enum):
    GILMORE          = 1
    KELLER_MIKSIS    = 2
    RAYLEIGH_PLESSET = 3


class ThermalModel(enum.Enum):
    ADIABATIC  = 1
    ISOTHERMAL = 2
    TRANSFER   = 3


class BoundaryCondition(enum.Enum):
    PERIODIC                           =  -1
    REFLECTIVE                         =  -2
    GHOST_CELL_EXTRAPOLATION           =  -3
    RIEMANN_EXTRAPOLATION              =  -4
    SLIP_WALL                          =  -5
    NON_REFLECTING_SUBSONIC_BUFFER     =  -6
    NON_REFLECTING_SUBSONIC_INFLOW     =  -7
    NON_REFLECTING_SUBSONIC_OUTFLOW    =  -8
    FORCE_FREE_SUBSONIC_OUTFLOW        =  -9
    CONSTANT_PRESSURE_SUBSONIC_OUTFLOW = -10
    SUPERSONIC_INFLOW                  = -11
    SUPERSONIC_OUTFLOW                 = -12


class DatabaseFormat(enum.Enum):
    SILO_HDF5 = 1
    BINARY    = 2


class FloatingPrecision(enum.Enum):
    SINGLE = 1
    DOUBLE = 2


class RiemannSolver(enum.Enum):
    HLL   = 1
    HLLC  = 2
    EXACT = 3


class WaveSpeedEstimation(enum.Enum):
    DIRECT            = 1
    PRESSURE_VELOCITY = 2


class AverageStateEvaluation(enum.Enum):
    ROE_MEAN        = 1
    ARITHMETIC_MEAN = 2


class TimeStepper(enum.Enum):
    RUNGE_KUTTA_1 = 1
    RUNGE_KUTTA_2 = 2
    RUNGE_KUTTA_3 = 3
    RUNGE_KUTTA_4 = 4
    RUNGE_KUTTA_5 = 5


class WenoVariables(enum.Enum):
    CONSERVATIVE = 1
    PRIMITIVE    = 2


# === DEFINE UTILITY CLASSES === #


@dataclasses.dataclass
class Point:
    x: float = dataclasses.field(default=0)
    y: float = dataclasses.field(default=0)
    z: float = dataclasses.field(default=0)


class Vector(Point):
    pass


# ===    DEFINE CLASSES    === #
# === COMPUTATIONAL DOMAIN === #


@dataclasses.dataclass
class AxisDomain:
    begin: float = dataclasses.field(default=0)
    end:   float = dataclasses.field(default=0)


@dataclasses.dataclass
class AxisStretch(AxisDomain):
    stretch: bool  = dataclasses.field(default=False)
    rate:    float = dataclasses.field(default=0)


@dataclasses.dataclass
class Stretch:
    x: AxisStretch = dataclasses.field(default=AxisStretch())
    y: AxisStretch = dataclasses.field(default=AxisStretch())
    z: AxisStretch = dataclasses.field(default=AxisStretch())


@dataclasses.dataclass
class Cells:
    x: int
    y: int = dataclasses.field(default=0)
    z: int = dataclasses.field(default=0)


@dataclasses.dataclass
class SpacialDomain:
    x: AxisDomain
    y: AxisDomain = dataclasses.field(default=AxisDomain(begin=0, end=0))
    z: AxisDomain = dataclasses.field(default=AxisDomain(begin=0, end=0))


@dataclasses.dataclass
class Time:
    end:   int
    save:  int
    dt:    float
    start: int   = dataclasses.field(default=0)


@dataclasses.dataclass
class ComputationalDomain:
    cells:        Cells
    domain:       SpacialDomain
    time:         Time
    cyl_coord:    bool    = dataclasses.field(default=False)
    stretch:      Stretch = dataclasses.field(default=Stretch())


# ===        FLUIDS        === #


@dataclasses.dataclass
class Fluid:
    gamma:   float      = None
    pi_inf:  float      = None
    mul0:    typing.Any = None
    ss:      typing.Any = None
    gamma_v: typing.Any = None
    M_v:     typing.Any = None
    mu_v:    typing.Any = None
    k_v:     typing.Any = None
    G:       typing.Any = None


# === SIMULATION ALGORITHM === #


@dataclasses.dataclass
class WenoParameters:
    order:     int
    variables: WenoVariables
    mapped:    bool
    mp:        bool
    epsilon:   float = dataclasses.field(default=1e-16)


@dataclasses.dataclass
class AxisBoundaryCondition:
    begin: BoundaryCondition
    end:   BoundaryCondition


@dataclasses.dataclass
class BoundaryConditions:
    x: AxisBoundaryCondition
    y: AxisBoundaryCondition = dataclasses.field(default=None)
    z: AxisBoundaryCondition = dataclasses.field(default=None)


@dataclasses.dataclass
class SimulationAlgorithm:
    model:          MulticomponentModel
    weno:           WenoParameters
    boundary:       BoundaryConditions
    riemann_solver: RiemannSolver
    time_stepper:   TimeStepper
    wave_speeds:    WaveSpeedEstimation
    avg_state:      AverageStateEvaluation
    mixture_err:    bool = dataclasses.field(default=False)
    mpp_lim:        bool = dataclasses.field(default=False)
    adv_alphan:     bool = dataclasses.field(default=False) 
    hypoelasticity: bool = dataclasses.field(default=False)
    null_weights:   bool = dataclasses.field(default=False)


# ===   PATCH PARAMETERS   === #


@dataclasses.dataclass
class Patch:
    geometry:        PatchGeometry
    radius:          float         = None
    radii:           typing.Any    = None
    epsilon:         typing.Any    = None
    beta:            typing.Any    = None
    smooth_patch_id: int           = None
    alpha:           typing.Any    = None
    alpha_rho:       typing.Any    = None
    smooth_coeff:    float         = None
    rho:             typing.Any    = None
    pressure:        float         = None
    gamma:           float         = None
    pi_inf:          float         = None
    r0:              typing.Any    = None
    v0:              typing.Any    = None
    p0:              typing.Any    = None
    m0:              typing.Any    = None
    centroid:        Point         = None
    length:          Vector        = None
    radii:           Vector        = None
    normal:          Vector        = None
    smoothen:        bool          = dataclasses.field(default=False)
    velocity:        Vector        = dataclasses.field(default=Vector(x=0,y=0,z=0))

    # TODO: Add for arho_id in range(1, 10+1):

    alter_patch:     typing.Any = dataclasses.field(default=False)


# ===  DATABASE STRUCTURE  === #


@dataclasses.dataclass
class DatabaseWrite:
    probe:      bool = dataclasses.field(default=False)
    cons_vars:  bool = dataclasses.field(default=False)
    prim_vars:  bool = dataclasses.field(default=False)
    rho:        bool = dataclasses.field(default=False)
    E:          bool = dataclasses.field(default=False)
    pres:       bool = dataclasses.field(default=False)
    gamma:      bool = dataclasses.field(default=False)
    heat_ratio: bool = dataclasses.field(default=False)
    pi_inf:     bool = dataclasses.field(default=False)
    pres_inf:   bool = dataclasses.field(default=False)
    c:          bool = dataclasses.field(default=False)
    schlieren:  bool = dataclasses.field(default=False)
    #TODO: Some (i) params


@dataclasses.dataclass
class DatabseStructure:
    precision:      FloatingPrecision 
    fd_order:       int
    write:          DatabaseWrite     = dataclasses.field(default=DatabaseWrite())
    alt_soundspeed: bool              = dataclasses.field(default=False)
    parallel_io:    bool              = dataclasses.field(default=False)
    format:         DatabaseFormat    = dataclasses.field(default=DatabaseFormat.SILO_HDF5)


# ===       BUBBLES       === #


@dataclasses.dataclass
class Bubbles:
    bubbles:    bool = None
    model:      BubbleModel = None
    polytropic: bool = None
    thermal:    ThermalModel = None
    R0ref:      float = None
    number:     int   = None
    cavitation: float = None
    weber:      float = None
    Re_inv:     float = None
    mu_10:      float = None
    ss:         float = None
    pv:         float = None
    gamma_v:    float = None
    M_v:        float = None
    mu_v:       float = None
    k_v:        float = None
    qbmm:       bool  = None


# === DEFINE A CASE === #


@dataclasses.dataclass
class Case:
    domain:    ComputationalDomain
    fluids:    typing.List[Fluid]
    patches:   typing.List[Patch]
    bubbles:   Bubbles
    database:  DatabseStructure
    algorithm: SimulationAlgorithm

    def __repr__(self) -> str:
        return json.dumps(self.to_json(), indent=4)

    def to_json(self) -> dict:
        def handle_enums(data):
            def convert_value(obj):
                if isinstance(obj, enum.Enum):
                    return obj.value

                return obj

            return dict((key, convert_value(value)) for key, value in data)

        return dataclasses.asdict(self, dict_factory=handle_enums)


# === DEFINE ALL ENUMERATIONS === # # === DEFINE UTILITY CLASSES === # # ===    DEFINE CLASSES    === # # === COMPUTATIONAL DOMAIN === # # ===        FLUIDS        === # # === SIMULATION ALGORITHM === # # ===   PATCH PARAMETERS   === # # TODO: Add for arho_id in range(1, 10+1): # ===  DATABASE STRUCTURE  === # #TODO: Some (i) params # ===       BUBBLES       === # # === DEFINE A CASE === #