import numpy as np
import pytest


from vpipe_flow_calc.extra.dummy_plotter import LivePlotter
from vpipe_flow_calc.solver.dummy_newton_solver import NewtonSolver
from vpipe_flow_calc.solver.time_stepper import TS

PLOT_RESULTS = True


def delete_results_from_temp_folder(folder):
    import os
    for the_file in os.listdir(folder):
        file_path = os.path.join(folder, the_file)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)

        except Exception as e:
            print(e)

R = 286.9 # Air individual gas constant [J/(kg.K)]       
def ideal_gas_density_model(P, T=288.15):        
    return P / (R * T)


c_L = 5e-11
def constant_density_model(P):
    return 1000 * (1 + c_L * (P - 1e5))


def liquid_viscosity_model(P):
    return np.full_like(P, 5e-2)


def gas_viscosity_model(P):
    return np.full_like(P, 5e-6)

density_model = [ideal_gas_density_model, constant_density_model, constant_density_model, ideal_gas_density_model]
viscosity_model = [gas_viscosity_model, liquid_viscosity_model, liquid_viscosity_model, gas_viscosity_model]

def surface_tension_model(P):
    return np.full_like(P, 7.197e-2)


N = 201

@pytest.fixture
def Δt():
    return 1e-4 # [s]

@pytest.fixture
def Δt_min():
    return 1e-12 # [s]

@pytest.fixture
def Δt_max():
    return 1.0 # [s]

@pytest.fixture
def pipe_length():
    return 20.0 # [m]

@pytest.fixture
def diameter():
    return 0.05 # [m]

@pytest.fixture
def P_presc():
    return 1.0 # [bar]

@pytest.fixture
def ṁ_presc():
    U_SL = 0.1
    U_SG = np.nan
    
    ρ_L = constant_density_model(1e5)
    ρ_G = ideal_gas_density_model(1e5)
    
    A = 0.25 * np.pi * diameter() * diameter()
    
    ṁ_L = U_SL * ρ_L * A
    ṁ_G = U_SG * ρ_G * A
    
    return np.array([0, 0.1, 0, np.nan]) # [kg/s]
#     return np.array([0, ṁ_L, 0, ṁ_G]) # [kg/s]
#     return np.array([0, 2.5, 0, 0.005]) # [kg/s]


@pytest.fixture
def θ():
    a = np.full(N, 0.5 * np.pi)
#     a[-1] = 0
    return a

@pytest.fixture
def ε():
    return np.full(N, 1e-5)


# @pytest.mark.skip(reason='not testing it now...')
@pytest.mark.timeout(0)
def test_dispersed_vertical(
    Δt, Δt_min, Δt_max,
    pipe_length, diameter, ε, θ,
    P_presc, ṁ_presc
):
#     suffix = '_'
    from vpipe_flow_calc.mass_equation import flux_limiter
    suffix = '_case_drift_fix_high_velocity_gas_full_cont_' + flux_limiter + '_sym'
    folder = '/home/soprano/work/vpipe-flow/temp/results/profile' + suffix

    import os
    if not os.path.exists(folder):
        os.makedirs(folder)

    delete_results_from_temp_folder(folder)


    n_phases = 2
    n_layers = 2
    n_fields = 4

    dof = n_fields + n_layers + 1

    U = np.full((n_layers, N), 1e-8)

    α = np.full((n_fields, N), np.nan)
    little_value = 1e-15
    α[0] = 0.001 - little_value
    α[1] = 0.99 - little_value
    α[2] = 0.0 + little_value
    α[3] = 0.01 + little_value

    P = np.full(N, P_presc)
    mass_error = np.zeros_like(α)


    log = lambda expr: print(expr)
    solver = NewtonSolver(N, dof, pipe_length, diameter, ε, θ, n_phases, n_fields, n_layers, ṁ_presc, P_presc, density_model, viscosity_model, surface_tension_model, log=log)
    ts = TS(solver, log=log)
    
    initial_time = 0
    time_intervals = np.linspace(0.1, 350, num=3500) # [s]

    initial_time = 30.0015741903
    time_intervals = np.linspace(30.1, 140.8, num=3100) # [s]
               
    filename = '/home/soprano/work/vpipe-flow/temp/results/profile_30.0015741903.npz'
    restart_conditions = np.load(filename)
    α = restart_conditions['α']
    P = restart_conditions['P']
    U = restart_conditions['U_layer']
    mass_error = restart_conditions['mass_error']
    
#     initial_time = 130.700279571
#     time_intervals = np.linspace(130.8, 140.8, num=100) # [s]
#               
#     filename = '/home/soprano/work/vpipe-flow/temp/results/profile_130.700279571.npz'
#     restart_conditions = np.load(filename)
#     α = restart_conditions['α']
#     P = restart_conditions['P']
#     U = restart_conditions['U_layer']
#     mass_error = restart_conditions['mass_error']
    
    if PLOT_RESULTS:
        probe_position = 10
        dx = pipe_length/(N-1)
        plotter = LivePlotter(N, pipe_length, dx, probe_position, P_presc, final_time=50)

    
    for i, final_time in enumerate(time_intervals):
        solver.set_initial_solution(U, α, P, mass_error)
        ts.set_duration(initial_time, final_time)
        ts.set_initial_timestep(Δt)
        ts.set_min_timestep(Δt_min)
        ts.set_max_timestep(Δt_max)
        ts.solve()

        U, α, P, _, _, _ = solver.get_UαP_array()

        initial_time = final_time
        Δt = ts.Δt

        if PLOT_RESULTS:
            solver.normalize_vol_frac(fix_range=True)              
            x = solver.x
            x_f = solver.x_f
#             plotter.plot_profile_results(α, P, U, solver.U_fields)
#             plotter.plot_trend_results(ts.current_time, α, P, U)
#             plotter.draw()
            plotter.save(x, x_f, α, P, U, solver.U_fields, solver.mass_error,
                         filename=folder + '/profile_' + str(ts.current_time))

#     positions = [2, 500, 900]
#
#     U, α, P, _, _, _ = solver.get_UαP_array()
#
#     obtained_U = U[:, positions]
#     obtained_α = α[:, positions]
#     obtained_P = P[positions]
#
#     assert approx(obtained_U, rtol=1e-4) == expected_U
#     assert approx(obtained_α, rtol=1e-4) == expected_α
#     assert approx(obtained_P, rtol=1e-4) == expected_P


@pytest.mark.skip(reason='not testing it now...')
@pytest.mark.timeout(0)
def test_entrainment(
    Δt, Δt_min, Δt_max,
    pipe_length, diameter, ε, θ,
    P_presc, ṁ_presc
):
    suffix = '_case_entraiment'
    folder = '/home/soprano/work/vpipe-flow/temp/results/profile' + suffix

    import os
    if not os.path.exists(folder):
        os.makedirs(folder)

    delete_results_from_temp_folder(folder)

    n_phases = 2
    n_layers = 2
    n_fields = 4

    dof = n_fields + n_layers + 1

    log = lambda expr: print(expr)
    solver = NewtonSolver(N, dof, pipe_length, diameter, ε, θ, n_phases, n_fields, n_layers, ṁ_presc, P_presc, density_model, viscosity_model, surface_tension_model, log=log)
    ts = TS(solver, log=log)
    
    initial_time = 310.000134864
    time_intervals = np.linspace(310.1, 350, num=400) # [s]
        
    filename = '/home/soprano/work/vpipe-flow/temp/results/profile_310.000134864.npz'
    restart_conditions = np.load(filename)
    α = restart_conditions['α']
    P = restart_conditions['P']
    U = restart_conditions['U_layer']
    mass_error = restart_conditions['mass_error']
    
    probe_position = 10
    dx = pipe_length/(N-1)
    plotter = LivePlotter(N, pipe_length, dx, probe_position, P_presc, final_time=50)
    
    for i, final_time in enumerate(time_intervals):
        solver.set_initial_solution(U, α, P, mass_error)
        ts.set_duration(initial_time, final_time)
        ts.set_initial_timestep(Δt)
        ts.set_min_timestep(Δt_min)
        ts.set_max_timestep(Δt_max)
        ts.solve()

        U, α, P, _, _, _ = solver.get_UαP_array()

        initial_time = final_time
        Δt = ts.Δt

        solver.normalize_vol_frac(fix_range=True)              
        x = solver.x
        x_f = solver.x_f
        plotter.save(x, x_f, α, P, U, solver.U_fields, solver.mass_error,
                     filename=folder + '/profile_' + str(ts.current_time))

