import os
import numpy as np
import json
import platform
from svv.simulation.fluid.rom.zero_d.post import make_results, view_plots, post_data, collate_timeseries_to_pvd
from svv.simulation.fluid.rom.zero_d.process import run_0d_script
from svv.simulation.fluid.inflow.waveform import generate_physiologic_wave
from svv.utils.solvers.solver_0d import get_solver_0d_exe

def export_0d_simulation(forest, network_id, inlets, steady=True, outdir=None, folder="0d_tmp", number_cardiac_cycles=1,
                            flow=None, number_time_pts_per_cycle=5,material="olufsen",
                            olufsen={'k1': 0.0, 'k2': -22.5267, 'k3': 1.0e7, 'material exponent': 2.0, 'material pressure': 0.0},
                            linear={'material ehr': 1e7, 'material pressure': 0.0}, get_0d_solver=False, path_to_0d_solver=None,
                            viscosity_model='constant', vivo=True, distal_pressure=0.0, capacitance=True, inductance=True):
    """
    Export 0D simulation.
    """
    """
    Method for exporting 0d simulations for generic forest
    objects.

    Parameters
    ----------
                network_id : int
                       the index of the network 
                steady : bool (default = True)
                       true/false flag for assigning a time-varying inlet
                       flow. If false, a flow waveform is assigned from
                       empirical data taken from coronary arteries.
                       Refer to svcco.sv_interface.waveform for more details
                outdir : str (default = None)
                       string specifying the output directory for the 0D simulation
                       folder to be created
                folder : str (default = "0d_tmp")
                       string specifying the folder name. This folder will store
                       all 0D simulation files that are generated
                number_cardiac_cycles : int (default = 1)
                       number of cardiac cycles to compute the 0D simulation.
                       This can be useful for obtaining a stable solution.
                number_time_pts_per_cycle : int (default = 5)
                       number of time points to evaluate during the 0D
                       simulation. A default of 5 is good for steady flow cases;
                       however, pulsatile flow will likely require a higher number
                       of time points to adequetly capture the inflow waveform.
                density : float (default = 1.06)
                       density of the fluid within the vascular tree to simulate.
                       By deafult, the density of blood is used.
                       Base Units are in centimeter-grame-second (CGS)
                viscosity : float (default = 0.04)
                       viscosity of the fluid within the vascular tree to simulate.
                       By default, the viscosity if blood is used.
                       Base units are in centimeter-grame-second (CGS)
                material : str (default = "olufsen")
                       string specifying the material model for the vascular
                       wall. Options available: ("olufsen" and "linear")
                olufsen : dict
                         material paramters for olufsen material model
                               'k1' : float (default = 0.0)
                               'k2' : float (default = -22.5267)
                               'k3' : float (default = 1.0e7)
                               'material exponent' : int (default = 2)
                               'material pressure' : float (default = 0.0)
                         By default these parameters numerically yield a "rigid"
                         material model.
                linear : dict
                         material parameters for the linear material model
                               'material ehr' : float (default = 1e7)
                               'material pressure' : float (default = 0.0)

                         By default these parameters numerically yield a "rigid"
                         material model.
                get_0d_solver : bool (default = False)
                         true/false flag to search for the 0D solver. If installed,
                         the solver will be imported into the autogenerated
                         simulation code. Depending on how large the disk-space
                         is for the current machine, this search may be prohibitively
                         expensive if no path is known.
                path_to_0d_solver : str (default = None)
                         string specifying the path to the 0D solver
                viscosity_model : str (default = 'constant')
                         string specifying the viscosity model of the fluid.
                         This may be useful for multiphase fluids (like blood)
                         which have different apparent viscosities at different
                         characteristic length scales. By default, constant
                         viscosity is used.

                         Other viscosity option: 'modified viscosity law'
                              This option leverages work done by Secomb et al.
                              specifically for blood in near-capillary length
                              scales.
                vivo : bool (default = True)
                         true/false flag for using in vivo blood viscosity hematocrit
                         value. If false, in vitro model is used. (Note these
                         are only used for the modified viscosity law model)
                distal_pressure : float (default = 0.0)
                         apply a uniform distal pressure along outlets of the
                         vascular tree
    Returns
    -------
                None
    """
    if outdir is None:
        outdir = os.path.join(os.getcwd(), folder)
    else:
        outdir = os.path.join(outdir, folder)

    # Ensure the target directory exists without failing if it already does.
    os.makedirs(outdir, exist_ok=True)
    resolved_solver_0d = None
    if path_to_0d_solver is not None:
        resolved_solver_0d = path_to_0d_solver
        print(f"Using explicit svZeroDSolver path: {resolved_solver_0d}")
    elif get_0d_solver:
        try:
            resolved_solver_0d = str(get_solver_0d_exe())
            print(f"Using svZeroDSolver executable: {resolved_solver_0d}")
        except Exception as e:
            print(f"WARNING: No svZeroDSolver executable found ({e})")
            print("Generated run.py will include setup instructions to install/build solver support.")
    input_file = {'description': {'description of case': None,
                                  'analytical results': None},
                  'boundary_conditions': [],
                  'junctions': [],
                  'simulation_parameters': {},
                  'vessels': []}
    simulation_parameters = {}
    simulation_parameters["number_of_cardiac_cycles"] = number_cardiac_cycles
    simulation_parameters["number_of_time_pts_per_cardiac_cycle"] = number_time_pts_per_cycle
    input_file['simulation_parameters'] = simulation_parameters

    def build_vessel_segment(pt_0, pt_1, r, idx, kinematic_viscosity, density, material=material,
                             viscosity_model=viscosity_model):
        vessel_segment = {}
        vessel_segment['vessel_id'] = idx
        vessel_segment['vessel_length'] = np.linalg.norm(pt_1 - pt_0)
        vessel_segment['vessel_name'] = 'branch' + str(idx) + '_seg0'
        vessel_segment['zero_d_element_type'] = 'BloodVessel'
        if material == 'olufsen':
            material_stiffness = olufsen['k1'] * np.exp(olufsen['k2'] * r) + olufsen['k3']
        else:
            material_stiffness = linear['material ehr']
        if viscosity_model == 'constant':
            nu = kinematic_viscosity * density
        elif viscosity_model == 'modified viscosity law':
            W = 1.1
            lam = 0.5
            D = r * 2 * (10000)
            nu_ref = kinematic_viscosity
            Hd = 0.45  # discharge hematocrit (dimension-less)
            C = lambda d: (0.8 + np.exp(-0.075)) * (-1 + (1 / (1 + 10 ** (-11) * d ** 12))) + (
                        1 / (1 + 10 ** (-11) * d ** 12))
            nu_vitro_45 = lambda d: 220 * np.exp(-1.3 * d) + 3.2 - 2.44 * np.exp(-0.06 * d ** 0.645)
            nu_vivo_45 = lambda d: 6 * np.exp(-0.085 * d) + 3.2 - 2.44 * np.exp(-0.06 * d ** 0.645)
            if vivo == True:
                nu_45 = nu_vivo_45
            else:
                nu_45 = nu_vitro_45
            nu_mod = lambda d: (1 + (nu_45 - 1) * (((1 - Hd) ** C - 1) / ((1 - 0.45) ** C - 1)) * (d / (d - W)) ** (
                        4 * lam)) * (d / (d - W)) ** (4 * (1 - lam))
            ref = nu_ref - nu_mod(10000)  # 1cm (units given in microns)
            nu = nu_mod(D) + ref
        zero_d_element_values = {}
        zero_d_element_values["R_poiseuille"] = ((8 * nu / np.pi) * vessel_segment['vessel_length']) / r ** 4
        if capacitance:
            zero_d_element_values["C"] = (3 * vessel_segment['vessel_length'] * np.pi * r ** 2) / (2 * material_stiffness)
        else:
            zero_d_element_values["C"] = 0.0
        if inductance:
            zero_d_element_values["L"] = (vessel_segment['vessel_length'] * density) / (np.pi * r ** 2)
        else:
            zero_d_element_values["L"] = 0.0
        zero_d_element_values["stenosis_coefficient"] = 0.0
        vessel_segment['zero_d_element_values'] = zero_d_element_values
        return vessel_segment

    vessels = {'models':[], 'proximal_points':[], 'distal_points':[], 'radii':[]}
    vessel_shift = 0
    bc_vessels = []
    # BUILD VESSELS FOR NETWORK
    for tree_id in range(forest.n_trees_per_network[network_id]):
        kinematic_viscosity = forest.networks[network_id][tree_id].parameters.kinematic_viscosity
        density = forest.networks[network_id][tree_id].parameters.fluid_density
        # BUILD TREE VESSEL SEGMENTS
        networks = forest.connections.tree_connections[network_id].connected_network
        for vessel_id in range(networks[tree_id].data.shape[0]):
            if tree_id in inlets:
                pt_0 = networks[tree_id].data[vessel_id, 0:3]
                pt_1 = networks[tree_id].data[vessel_id, 3:6]
            else:
                pt_0 = networks[tree_id].data[vessel_id, 3:6]
                pt_1 = networks[tree_id].data[vessel_id, 0:3]
            r = networks[tree_id].data[vessel_id, 21]
            vessel_segment = build_vessel_segment(pt_0, pt_1, r, len(vessels['models']), kinematic_viscosity,
                                                  density, material, viscosity_model)
            if vessel_id == 0:
                if tree_id in inlets:
                    bc = {}
                    bc['bc_name'] = "INFLOW"
                    bc['bc_type'] = "FLOW"
                    bc_values = {}
                    if steady:
                        q_in = float(networks[tree_id].data[vessel_id, 22]) if isinstance(flow, type(None)) else float(flow)
                        bc_values["Q"] = [q_in, q_in]
                        bc_values["t"] = [0, 1]
                    else:
                        if isinstance(flow, type(None)):
                            time, flow_wave = generate_physiologic_wave(
                                networks[tree_id].data[vessel_id, 22],
                                networks[tree_id].data[vessel_id, 21] * 2,
                            )
                            bc_values["Q"] = np.asarray(flow_wave, dtype=float).tolist()
                            bc_values["t"] = np.asarray(time, dtype=float).tolist()
                        else:
                            q_in = float(flow)
                            bc_values["Q"] = [q_in, q_in]
                            bc_values["t"] = [0, 1]
                        if bc_values["Q"]:
                            bc_values["Q"][-1] = bc_values["Q"][0]
                        simulation_parameters["number_of_time_pts_per_cardiac_cycle"] = len(bc_values["Q"])
                    with open(outdir + os.sep + "inflow.flow", "w") as file:
                        for i in range(len(bc_values["t"])):
                            file.write("{}  {}\n".format(bc_values["t"][i], bc_values["Q"][i]))
                    file.close()
                    bc['bc_values'] = bc_values
                    input_file['boundary_conditions'].append(bc)
                    vessel_segment['boundary_conditions'] = {'inlet': 'INFLOW'}
                    bc_vessels.append(len(vessels['models']))
                else:
                    bc = {}
                    bc['bc_name'] = 'OUT' + str(len(vessels['models']))
                    bc['bc_type'] = "RESISTANCE"
                    bc_values = {}
                    bc_values["Pd"] = distal_pressure  # self.parameters["Pterm"]
                    bc_values["R"] = 0  # total_resistance*(total_outlet_area/(np.pi*self.networks[network_id][tree_id].data[vessel_id,21]**2))
                    bc['bc_values'] = bc_values
                    input_file['boundary_conditions'].append(bc)
                    vessel_segment['boundary_conditions'] = {'outlet': 'OUT' + str(len(vessels['models']))}
                    bc_vessels.append(len(vessels['models']))
            vessels['models'].append(vessel_segment)
            vessels['proximal_points'].append(pt_0)
            vessels['distal_points'].append(pt_1)
            vessels['radii'].append(r)
            #input_file['vessels'].append(vessel_segment)
        # BUILD CONNECTING VESSEL SEGMENTS
        for conn_id in range(len(forest.connections.tree_connections[network_id].vessels[tree_id])):
            for vessel_id in range(forest.connections.tree_connections[network_id].vessels[tree_id][conn_id].data.shape[0]):
                if tree_id in inlets:
                    pt_0 = forest.connections.tree_connections[network_id].vessels[tree_id][conn_id][vessel_id, 0:3]
                    pt_1 = forest.connections.tree_connections[network_id].vessels[tree_id][conn_id][vessel_id, 3:6]
                else:
                    pt_0 = forest.connections.tree_connections[network_id].vessels[tree_id][conn_id][vessel_id, 3:6]
                    pt_1 = forest.connections.tree_connections[network_id].vessels[tree_id][conn_id][vessel_id, 0:3]
                r = forest.connections.tree_connections[network_id].vessels[tree_id][conn_id][vessel_id, 6]
                vessel_segment = build_vessel_segment(pt_0, pt_1, r, len(vessels['models']),
                                                      kinematic_viscosity, density, material, viscosity_model)
                vessels['models'].append(vessel_segment)
                vessels['proximal_points'].append(pt_0)
                vessels['distal_points'].append(pt_1)
                vessels['radii'].append(r)
    vessels['proximal_points'] = np.array(vessels['proximal_points'])
    vessels['distal_points'] = np.array(vessels['distal_points'])
    vessels['radii'] = np.array(vessels['radii'])
    # BUILD JUNCTIONS FOR TREES
    junctions = []
    points = np.unique(np.vstack((vessels['proximal_points'], vessels['distal_points'])), axis=0)
    for j_id in range(points.shape[0]):
        junction = {}
        junction['junction_name'] = "J" + str(len(junctions))
        junction['junction_type'] = "NORMAL_JUNCTION"
        inlet_vessels = np.argwhere(np.all(vessels['distal_points'] == points[j_id, :], axis=1)).flatten()
        outlet_vessels = np.argwhere(np.all(vessels['proximal_points'] == points[j_id, :], axis=1)).flatten()
        if inlet_vessels.shape[0] > 0:
            junction['inlet_vessels'] = inlet_vessels.tolist()
        else:
            continue
        if outlet_vessels.shape[0] > 0:
            junction['outlet_vessels'] = outlet_vessels.tolist()
        else:
            continue
        junctions.append(junction)

    # ADD VESSELS AND JUNCTIONS TO INPUT FILE
    input_file['vessels'] = vessels['models']
    input_file['junctions'] = junctions


    obj = json.dumps(input_file, indent=4)
    with open(outdir + os.sep + "solver_0d.in", "w") as file:
        file.write(obj)
    file.close()

    with open(outdir + os.sep + "plot_0d_results_to_3d.py", "w") as file:
        file.write(make_results)
    file.close()

    with open(outdir + os.sep + "plot_0d_results_at_slices.py", "w") as file:
        file.write(view_plots)
    file.close()

    with open(outdir + os.sep + "post_process.py", "w") as file:
        file.write(post_data)
    file.close()

    with open(outdir + os.sep + "collate_timeseries_to_pvd.py", "w") as file:
        file.write(collate_timeseries_to_pvd)
    file.close()

    with open(outdir + os.sep + "run.py", "w") as file:
        if resolved_solver_0d is not None:
            cmd_file = outdir + os.sep + ("cmd_run.bat" if platform.system() == "Windows" else "cmd_run.bash")
            with open(cmd_file, "w") as cmd:
                cmd.write(f"\"{resolved_solver_0d}\" \"solver_0d.in\"")
        file.write(run_0d_script.format(solver_exe=resolved_solver_0d, input_filename="solver_0d.in"))
    file.close()
    # DETERMINE NUMBER OF ROWS FOR GEOM
    geoms = np.zeros((vessels['proximal_points'].shape[0], 8))
    geoms[:, 0:3] = vessels['proximal_points']
    geoms[:, 3:6] = vessels['distal_points']
    geoms[:, 6] = np.linalg.norm(vessels['distal_points'] - vessels['proximal_points'], axis=1).flatten()
    geoms[:, 7] = vessels['radii']

    np.savetxt(outdir + os.sep + "geom.csv", geoms, delimiter=",")
