import numpy as np
import os
import json
import platform

def export_0d_simulation(tree ,steady=True ,outdir=None ,folder="0d_tmp" ,number_cardiac_cycles=1,flow=None,
                         number_time_pts_per_cycle=5 ,density=1.06 ,viscosity=0.04 ,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,
                         capacitance=True, inductance=True, filename="solver_0d.in", geom_filename="geom.csv"):
    """
    This script builds the 0D input file for running 0D simulation.

    Parameters:
    -----------
                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.getcwd() +os.sep +folder
    else:
        outdir = outdir +os.sep +folder
    if not os.path.isdir(outdir):
        os.mkdir(outdir)
    if get_0d_solver:
        #if path_to_0d_solver is None:
        #    path_to_0d_solver = locate_0d_solver()
        #else:
        #    path_to_0d_solver = locate_0d_solver(windows_drive=path_to_0d_solver ,linux_drive=path_to_0d_solver)
        path_to_0d_solver = None
    else:
        path_to_0d_solver = None
    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
    terminal_vessels = np.argwhere(tree.data[: ,16 ]==-1).flatten()
    total_outlet_area = np.sum(np.pi *tree.data[terminal_vessels ,21 ]**2)
    total_resistance  = (tree.parameters.terminal_pressure - 1333 * distal_pressure ) /tree.data[0 ,22]
    for vessel in range(tree.data.shape[0]):
        tmp = {}
        tmp['vessel_id'] = vessel
        tmp['vessel_length'] = tree.data[vessel ,20]
        tmp['vessel_name'] = 'branch ' +str(vessel ) +"_seg0"
        tmp['zero_d_element_type'] = "BloodVessel"
        if material == "olufsen":
            material_stiffness = olufsen['k1' ] *np.exp(olufsen['k2' ] *tree.data[vessel ,21] ) +olufsen['k3']
        else:
            material_stiffness = linear['material ehr']
        zero_d_element_values = {}
        if viscosity_model == 'constant':
            nu = tree.parameters.kinematic_viscosity
        elif viscosity_model == 'modified viscosity law':
            W  = 1.1
            lam = 0.5
            D  = tree.data[vessel ,21 ] * 2 *(10000)
            nu_ref = tree.parameters.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["R_poiseuille"] = ((8 * nu / np.pi) * tree.data[vessel, 20]) / tree.data[vessel, 21] ** 4
        if capacitance:
            zero_d_element_values["C"] = (3 * tree.data[vessel, 20] * np.pi * tree.data[vessel, 21] ** 2) / (
                        2 * material_stiffness)
        else:
            zero_d_element_values["C"] = 0.0
        if inductance:
            zero_d_element_values["L"] = (tree.data[vessel, 20] * density) / (np.pi * tree.data[vessel, 21] ** 2)
        else:
            zero_d_element_values["L"] = 0.0
        zero_d_element_values["stenosis_coefficient"] = 0.0
        tmp['zero_d_element_values'] = zero_d_element_values
        if vessel == 0:
            bc = {}
            bc['bc_name'] = "INFLOW"
            bc['bc_type'] = "FLOW"
            bc_values = {}
            if steady:
                if isinstance(flow, type(None)):
                    bc_values["Q"] = [float(tree.data[vessel, 22]), float(tree.data[vessel, 22])]
                    bc_values["t"] = [0, 1]
                else:
                    bc_values["Q"] = [flow, flow]
                    bc_values["t"] = [0, 1]
                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()
            #else:
            #    time, flow = wave(tree.data[vessel, 22], tree.data[vessel, 21] * 2)  # changed wave function
            #    bc_values["Q"] = flow.tolist()
            #    bc_values["t"] = time.tolist()
            #    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)
            tmp['boundary_conditions'] = {'inlet': "INFLOW"}
            if not np.isnan(tree.data[vessel, 15]) and not np.isnan(tree.data[vessel, 16]):
                junction = {}
                junction['inlet_vessels'] = [vessel]
                junction['junction_name'] = "J" + str(vessel)
                junction['junction_type'] = "NORMAL_JUNCTION"
                junction['outlet_vessels'] = [int(tree.data[vessel, 15]), int(tree.data[vessel, 16])]
                input_file['junctions'].append(junction)
        elif np.isnan(tree.data[vessel, 15]) and np.isnan(tree.data[vessel, 16]):
            bc = {}
            bc['bc_name'] = "OUT" + str(vessel)
            bc['bc_type'] = "RESISTANCE"
            bc_values = {}
            bc_values["Pd"] = 0  # tree.parameters["Pterm"]
            bc_values["R"] = float(total_resistance * (total_outlet_area / (np.pi * tree.data[vessel, 21] ** 2)))
            bc['bc_values'] = bc_values
            input_file['boundary_conditions'].append(bc)
            tmp['boundary_conditions'] = {'outlet': 'OUT' + str(vessel)}
        else:
            junction = {}
            junction['inlet_vessels'] = [vessel]
            junction['junction_name'] = "J" + str(vessel)
            junction['junction_type'] = "NORMAL_JUNCTION"
            junction['outlet_vessels'] = [int(tree.data[vessel, 15]), int(tree.data[vessel, 16])]
            input_file['junctions'].append(junction)
        input_file['vessels'].append(tmp)
    obj = json.dumps(input_file, indent=4)
    with open(outdir + os.sep + filename, "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 + "run.py", "w") as file:
        if platform.system() == "Windows":
            if path_to_0d_solver is not None:
                solver_path = path_to_0d_solver.replace(os.sep, os.sep + os.sep)
            else:
                solver_path = path_to_0d_solver
                print("WARNING: Solver location will have to be given manually")
                print("Current solver path is: {}".format(solver_path))
            solver_file = (outdir + os.sep + filename).replace(os.sep, os.sep + os.sep)
        else:
            if path_to_0d_solver is not None:
                solver_path = path_to_0d_solver
            else:
                solver_path = path_to_0d_solver
                print("WARNING: Solver location will have to be given manually")
                print("Current solver path is: {}".format(solver_path))
            solver_file = outdir + os.sep + "solver_0d.in"
        #file.write(run_0d_script.format(solver_path, solver_file))
    file.close()

    geom = np.zeros((tree.data.shape[0], 8))
    geom[:, 0:3] = tree.data[:, 0:3]
    geom[:, 3:6] = tree.data[:, 3:6]
    geom[:, 6] = tree.data[:, 20]
    geom[:, 7] = tree.data[:, 21]
    np.savetxt(outdir + os.sep + geom_filename, geom, delimiter=",")
