import math

from warnings import warn
import numpy as np

from openalea.mtg import traversal

from openalea.hydroroot.length import fit_law
from openalea.hydroroot import radius, flux, conductance 
from openalea.hydroroot.generator import markov, measured_root # 21-12-14: FB __init__.py in src not doing job
from openalea.hydroroot.water_solute_transport import pressure_calculation_no_non_permeating_solutes, init_some_MTG_properties, pressure_calculation

def hydroroot_mtg(
    primary_length=0.15,
    delta=2.e-3,
    beta=0.25,
    order_max=5,
    segment_length=1e-4,
    nude_length=0.02,
    seed=2,
    ref_radius=1e-4,
    order_decrease_factor=0.7,
    length_data=None,
    n=None,
    **kwds
):
    """Simulate and generate a root system.

    :param primary_length: Float (Default value = 0.15)
    :param delta: Float (Default value = 2.e-3)
    :param beta: Float (Default value = 0.25)
    :param order_max: int (Default value = 5)
    :param segment_length: Float (Default value = 1e-4)
    :param nude_length: Float (Default value = 0.02)
    :param seed: int (Default value = 2)
    :param ref_radius: Float (Default value = 1e-4)
    :param order_decrease_factor: Float (Default value = 0.7)
    :param length_data: pandas dataframe (Default value = None)
    :param n: int (Default value = None)

    :returns:
        - g (MTG)
        - surface (float)
        - volume (float)

    """
    # F. Bauget 2022-08-12: added if-else to be able to use the function without length data, useful for usage demo
    if length_data:
        xl, yl = length_data
        length_law = fit_law(xl, yl, scale=segment_length)
    else:
        length_law = None

    # compute the architecture
    nb_nude_vertices = int(nude_length / segment_length)
    branching_delay = int(delta / segment_length)

    if n is None:
        nb_vertices = int(primary_length / segment_length)
    else:
        nb_vertices = n
        warn("Use primary_length instead")

    g = markov.markov_binary_tree(
        nb_vertices=nb_vertices,
        branching_variability=beta,
        branching_delay=branching_delay,
        length_law=length_law,
        nude_tip_length=nb_nude_vertices,
        order_max=order_max,
        seed=seed)

    # compute radius property on MTG
    g = radius.ordered_radius(g, ref_radius=ref_radius, order_decrease_factor=order_decrease_factor)

    # compute length property and parametrisation
    g = radius.compute_length(g, segment_length)
    g = radius.compute_relative_position(g)

    g, surface = radius.compute_surface(g)
    g, volume = radius.compute_volume(g)

    return g, surface, volume

def hydroroot_flow(
    g,
    segment_length=1.e-4,
    k0=300,
    Jv=0.1,
    psi_e=0.4,
    psi_base=0.1,
    axial_conductivity_data=None,
    radial_conductivity_data=None,
):
    r"""Flux and equivalent conductance calculation

    :param g: MTG
    :param segment_length: (float) - not used vertices length in hydroroot in m (Default value = 1.0e-4) #TODO not used delete it
    :param k0: (float) - not used radial conductivity in :math:`10^{-9}\ m.s^{-1}.MPa^{-1}` (Default value = 300) #TODO not used delete it
    :param Jv: (Float) not used because invert_model=True in Flux.Flux (Default value = 0.1) #TODO delete it or add invert_model in arguments
    :param psi_e: (Float) external hydrostatic potential in MPa (Default value = 0.4)
    :param psi_base: (Float)  root base hydrostatic potential in MPa (Default value = 0.1)
    :param axial_conductivity_data: (2 list of Float) axial conductance (:math:`10^{-9}\ m^4.MPa^{-1}.s^{-1}`) versus distance to tip (m) (Default value = None)
    :param radial_conductivity_data: (2 list of Float) radial conductivity (:math:`10^{-9}\ m.MPa^{-1}.s^{-1}`) versus distance to tip (m) (Default value = None)
    :returns:
        - g (MTG): the MTG with the following properties filled: K (axial conductance), k (radial donductivity),
    	        j (radial flux), J_out (axial flux), psi_in and psi_out (hydrostatic pressure into the root at the
    	        input and output of a MTG node
    	- Keq (float): the equivalent conductance of the whole root
    	- Jv_global (float): the outgoing flux at the root base

    """
    if axial_conductivity_data:
        # Compute K using axial conductance data
        xa, ya = axial_conductivity_data
        # commented line below, BUG correction, the global flux was diverging when decreasing segment_length
        # ya was not supposed to be multiplied by (segment_length / 1.e-4)
        # ya = list(np.array(ya) * (segment_length / 1.e-4))
        axial_conductivity_law = fit_law(xa, ya)

        g = conductance.fit_property_from_spline(g, axial_conductivity_law, 'position', 'K_exp')
        g = conductance.compute_K(g) # Fabrice 2020-01-17: calculation of K in dimension [L^3 P^(-1) T^(-1)]

    if radial_conductivity_data:
        xr, yr = radial_conductivity_data
        radial_conductivity_law = fit_law(xr, yr)

        g = conductance.fit_property_from_spline(g, radial_conductivity_law, 'position', 'k0')
        g = conductance.compute_k(g, k0='k0')

    # Compute the flux
    g = flux.flux(g, Jv, psi_e, psi_base, invert_model=True)

    Keqs = g.property('Keq')
    v_base = next(g.component_roots_at_scale_iter(g.root, scale=1))

    Keq = Keqs[v_base]
    if psi_e is None:
        Peq = g.property('Peq')
        Jv_global = Keq * (Peq[v_base] - psi_base)
    else:
        Jv_global = Keq * (psi_e - psi_base)

    return g, Keq, Jv_global

def hydroroot_solute_flow(
    g,
    psi_e=0.4,
    psi_base=0.1,
    k0=None,
    axial_conductivity_data=None,
    radial_conductivity_data=None,
    J_s = 1.0e-7, Ps = 1.0e-9, Cse = 13.96e-9, Ce = 0.0, sigma=1.0, Temp=298, C_base = None, eps = 1.0e-9):
    r"""Water and solute fluxes and equivalent conductance calculation on an MTG

    :param g: MTG
    :param psi_e: (Float) external hydrostatic potential in MPa (Default value = 0.4)
    :param psi_base: (Float)  root base hydrostatic potential in MPa (Default value = 0.1)
    :param k0: (float) - uniform radial conductivity in :math:`10^{-9}\ m.s^{-1}.MPa^{-1}` (Default value = None)
    :param axial_conductivity_data: (2 list of Float) axial conductance profile (:math:`10^{-9}\ m^4.MPa^{-1}.s^{-1}`) versus distance to tip (m) (Default value = None)
    :param radial_conductivity_data: (2 list of Float) radial conductivity profile (:math:`10^{-9}\ m.MPa^{-1}.s^{-1}`) versus distance to tip (m) (Default value = None)
    :param J_s: (float) active pumping rate in mol/(m2.s) (Default value = 1e-7)
    :param Ps: (float) permeability coefficient in m/s (Default value = 1e-9)
    :param Cse: (float) initial permeating solute concentration inside the root in mol/microL (Default value = 13.96e-9)
    :param Ce: (float) initial non-permeating solute concentration outside the root in mol/microL (Default value = 0.)
    :param sigma: (float) - the reflexion coefficient between 0 and 1 (Default value = 1.0)
    :param Temp: (float) - temperature in Kelvin (Default value = 298)
    :param C_base: (float) solute concentration at the root base in mol/microL (Default value = None)
    :param eps: (float) convergence criterium (Default value = 1e-9)
    :returns:
        - g (MTG): the MTG with the following properties filled: K (axial conductance), k (radial donductivity),
    	        j (radial flux), J_out (axial flux), psi_in and psi_out (hydrostatic pressure into the root at the
    	        input and output of a MTG node)
    	- Jv_global (float): the outgoing flux at the root base

    """
    xa, ya = axial_conductivity_data
    axial_conductivity_law = fit_law(xa, ya)

    # Compute K using axial conductance data
    g = conductance.fit_property_from_spline(g, axial_conductivity_law, 'position', 'K_exp')
    g = conductance.compute_K(g)

    if k0:
        g = conductance.setting_k0_according_to_order(g, k0, k0)
    else:
        xr, yr = radial_conductivity_data
        radial_conductivity_law = fit_law(xr, yr)
        g = conductance.fit_property_from_spline(g, radial_conductivity_law, 'position', 'k0')

    g = conductance.compute_k(g, k0='k0')

    g = flux.flux(g, psi_e = psi_e, psi_base = psi_base, invert_model=True)
    g = init_some_MTG_properties(g, tau=J_s, Cini=Cse, Cpeg_ini = Ce, t = 1, Ps = Ps)

    if Ce <= 0.0:
        calculation = pressure_calculation_no_non_permeating_solutes
    else:
        calculation = pressure_calculation

    # Newton-Raphson loop
    nb_v = g.nb_vertices()
    Fdx = 1.0
    Fdx_old = 1.
    while Fdx > eps:
        g, dx, data, row, col = calculation(g, Temp=Temp, sigma=sigma, Ce=Ce, Cse=Cse, Pe=psi_e,
                                            Pbase=psi_base, C_base=C_base)
        Fdx = math.sqrt(sum(dx ** 2.0)) / nb_v
        if abs(Fdx - Fdx_old) < eps: break
        Fdx_old = Fdx

    Jv_global = g.property('J_out')[1]

    return g, Jv_global

def hydroroot(
    primary_length=0.15,
    delta=2.e-3,
    beta=0.25,
    order_max=5,
    segment_length=1e-4,
    nude_length=0.02,
    seed=2,
    ref_radius=1e-4,
    order_decrease_factor=0.7,
    k0=300,
    Jv=0.1,
    psi_e=0.4,
    psi_base=0.1,
    length_data=None,
    axial_conductivity_data=None,
    radial_conductivity_data=None,
    n=None
):
    """Simulate a root system and compute global conductance and flux.
    see :func:`~main.hydroroot_mtg` and  :func:`~main.hydroroot_flow`
    deprecated
    :param primary_length:  (Default value = 0.15)
    :param delta:  (Default value = 2.e-3)
    :param beta:  (Default value = 0.25)
    :param order_max:  (Default value = 5)
    :param segment_length:  (Default value = 1e-4)
    :param nude_length:  (Default value = 0.02)
    :param seed:  (Default value = 2)
    :param ref_radius:  (Default value = 1e-4)
    :param order_decrease_factor:  (Default value = 0.7)
    :param k0:  (Default value = 300)
    :param Jv:  (Default value = 0.1)
    :param psi_e:  (Default value = 0.4)
    :param psi_base:  (Default value = 0.1)
    :param length_data:  (Default value = None)
    :param axial_conductivity_data:  (Default value = None)
    :param radial_conductivity_data:  (Default value = None)
    :param n:  (Default value = None)

    """
    g, surface, volume = hydroroot_mtg(primary_length=primary_length,
                                       delta=delta,
                                       beta=beta,
                                       order_max=order_max,
                                       segment_length=segment_length,
                                       nude_length=nude_length,
                                       seed=seed,
                                       ref_radius=ref_radius,
                                       order_decrease_factor=order_decrease_factor,
                                       length_data=length_data,
                                       n=n,
                                       )
    xa, ya = axial_conductivity_data
    # commented line below, BUG correction, the global flux was diverging when decreasing segment_length
    # ya was not supposed to be multiplied by (segment_length / 1.e-4)
    # ya = list(np.array(ya) * (segment_length / 1.e-4))
    axial_conductivity_law = fit_law(xa, ya)

    xr, yr = radial_conductivity_data
    radial_conductivity_law = fit_law(xr, yr)

    # Compute K using axial conductance data
    g = conductance.fit_property_from_spline(g, axial_conductivity_law, 'position', 'K_exp')
    g = conductance.compute_K(g)  # Fabrice 2020-01-17: calculation of K in dimension [L^3 P^(-1) T^(-1)]
    # Compute the flux

    g = conductance.fit_property_from_spline(g, radial_conductivity_law, 'position', 'k0')
    g = conductance.compute_k(g, k0='k0')

    # TODO: return Keq base and Jv
    g = flux.flux(g, Jv, psi_e, psi_base, invert_model=True)

    Keqs = g.property('Keq')
    v_base = next(g.component_roots_at_scale_iter(g.root, scale=1))

    Keq = Keqs[v_base]
    Jv_global = Keq * (psi_e - psi_base)

    return g, surface, volume, Keq, Jv_global


def hydroroot_from_data(
    primary_length=0.15,
    delta=2.e-3,
    beta=0.25,
    order_max=5,
    segment_length=1e-4,
    nude_length = 0.02,
    seed = 2,
    ref_radius = 1e-4,
    order_decrease_factor = 0.7,
    k0 = 300,
    Jv = 0.1,
    psi_e = 0.4,
    psi_base = 0.1,
    length_data=None,
    axial_conductivity_data=None,
    radial_conductivity_data=None,
    primary_length_data=None,
    lateral_length_data=None,
    ):
    """Reconstruct a root system and compute global conductance and flux.
    deprecated
    :param primary_length_data: list Float (Default value = None)
    :param lateral_length_data: list Float (Default value = None)
    :param return: 
    :param g: MTG
    :param surface: float
    :param volume: float
    :param Keq: float
    :param Jv_global: float
    :param see: func
    :param primary_length:  (Default value = 0.15)
    :param delta:  (Default value = 2.e-3)
    :param beta:  (Default value = 0.25)
    :param order_max:  (Default value = 5)
    :param segment_length:  (Default value = 1e-4)
    :param nude_length:  (Default value = 0.02)
    :param seed:  (Default value = 2)
    :param ref_radius:  (Default value = 1e-4)
    :param order_decrease_factor:  (Default value = 0.7)
    :param k0:  (Default value = 300)
    :param Jv:  (Default value = 0.1)
    :param psi_e:  (Default value = 0.4)
    :param psi_base:  (Default value = 0.1)
    :param length_data:  (Default value = None)
    :param axial_conductivity_data:  (Default value = None)
    :param radial_conductivity_data:  (Default value = None)

    """
    xl, yl = length_data
    length_law = fit_law(xl, yl, scale=segment_length)

    xa, ya = axial_conductivity_data
    # commented line below, BUG correction, the global flux was diverging when decreasing segment_length
    # ya was not supposed to be multiplied by (segment_length / 1.e-4)
    # ya = list(np.array(ya) * (segment_length / 1.e-4))
    axial_conductivity_law = fit_law(xa, ya)

    xr, yr = radial_conductivity_data
    radial_conductivity_law = fit_law(xr, yr)

    # compute the architecture
    nb_nude_vertices = int(nude_length / segment_length)
    branching_delay = int(delta / segment_length)

    g = measured_root.mtg_builder(
        primary_length,
        primary_length_data,
        lateral_length_data,
        segment_length=segment_length,
        branching_variability=beta,
        branching_delay=branching_delay,
        length_law=length_law,
        nude_tip_length=nb_nude_vertices,
        order_max=order_max,
        seed=seed)


    # compute radius property on MTG
    g = radius.ordered_radius(g, ref_radius=ref_radius, order_decrease_factor=order_decrease_factor)

    # compute length property and parametrisation
    g = radius.compute_length(g, segment_length)
    g = radius.compute_relative_position(g)

    # Compute K using axial conductance data
    g = conductance.fit_property_from_spline(g, axial_conductivity_law, 'position', 'K_exp')
    g = conductance.compute_K(g)  # Fabrice 2020-01-17: calculation of K in dimension [L^3 P^(-1) T^(-1)]

    g, surface = radius.compute_surface(g)
    g, volume = radius.compute_volume(g)

    # Compute the flux

    g = conductance.fit_property_from_spline(g, radial_conductivity_law, 'position', 'k0')
    g = conductance.compute_k(g, k0='k0')

    # TODO: return Keq base and Jv
    g = flux.flux(g, Jv, psi_e, psi_base, invert_model=True)

    Keqs = g.property('Keq')
    v_base = next(g.component_roots_at_scale_iter(g.root, scale=1))

    Keq = Keqs[v_base]
    Jv_global = Keq * (psi_e - psi_base)

    return g, surface, volume, Keq, Jv_global

def root_builder(primary_length = 0.13, seed = None, delta = 2.0e-3, nude_length = 2.0e-2, df = None, segment_length = 1.0e-4,
                  length_data = None, branching_variability = 0.25, order_max = 4.0, order_decrease_factor = 0.7,
                  ref_radius = 7.0e-5, Flag_radius = False):
    """wrapper function: build a MTG with properties that are set like : radius, vertex length, position (distance to tip) and mylength (distance to base).
    The MTG is either generated or built from a DataFrame. The unit length in hydroroot should be in m.

    :param primary_length: (float) - primary root length for generated mtg (Default value = 0.13)
    :param seed: (int) - seed for generated mtg (Default value = None)
    :param delta: (float) - branching delay for generated mtg in hydroroot should be in m (Default value = 2.0e-3)
    :param nude_length: (float) - length from tip without lateral for generated mtg (Default value = 2.0e-2)
    :param df: (DataFrame) - DataFrame with the architecture data to be reconstructed if not None (Default value = None)
    :param segment_length: (float) - vertices length in hydroroot should be in m (Default value = 1.0e-4)
    :param length_data: (string list) - the file name with the length data laws (Default value = None)
    :param branching_variability: (float) (Default value = 0.25)
    :param order_max: (float) (Default value = 4.0)
    :param order_decrease_factor: (float) (Default value = 0.7)
    :param ref_radius: (float) - radius in hydroroot should be in m (Default value = 7.0e-5)
    :param Flag_radius: (boolean) - True if radius exist in df (Default value = False)
    :returns:
        - g: MTG with the different properties set or computed (see comments above),
        - primary_length: primary root length (output for generated mtg)
        - total_length: total root length
        - surface: total root surface
        - _seed: the seed used in the generator

    see also :func:`generator.measured_root.mtg_from_aqua_data` for details about df
    """
    if df is not None:
        g = measured_root.mtg_from_aqua_data(df, segment_length)
        _seed = None
    else:
        # if no seed just create one
        if seed is None:
            _seed = markov.my_seed()
        else:
            _seed = seed

        g = markov.generate_g(_seed, length_data,
                       branching_variability, delta,
                       nude_length, primary_length, segment_length,
                       order_max)

    # compute radius property on MTG
    # F. Bauget 2022-05-17 : added if
    if (not Flag_radius) or ('radius' not in g.property_names()):
        g = radius.ordered_radius(g, ref_radius, order_decrease_factor)

    # compute length property and parametrisation
    g = radius.compute_length(g, segment_length)
    g = radius.compute_relative_position(g)

    # Calculation of the distance from base of each vertex, used for cut and flow
    _mylength = {}
    for v in traversal.pre_order2(g, 1):
        pid = g.parent(v)
        _mylength[v] = _mylength[pid] + segment_length if pid else segment_length
    g.properties()['mylength'] = _mylength

    # total_length is the total length of the RSA (sum of the length of all the segments)
    total_length = g.nb_vertices(scale = 1) * segment_length
    g, surface = radius.compute_surface(g)
    g, volume = radius.compute_volume(g)

    if df is not None:
        v_base = next(g.component_roots_at_scale_iter(g.root, scale = g.max_scale()))
        primary_length = g.property('position')[v_base]

    return g, primary_length, total_length, surface, _seed

def soil_1D(g, soil_data = None):
    """
    add a soil as heterogeneous water potential along

    This is $\Psi_e$ versus depth with $depth>0$.

    :param g: MTG
    :param soil_data: tuple of 2 lists, (z,psi_e) z=depth (m), psi_e=water potential (MPa)
    :return: g
    """

    z,psi_e = soil_data

    # Compute absolute z coordinate and normalize
    vids = g.property('xyz').keys()
    zs = np.array([np.abs(pt.z) for pt in g.property('xyz').values()])
    zs = zs.tolist()

    # Fit data on z coordinate to compute psi_e on each vertex
    g.properties()['height'] = dict(zip(vids,zs))
    soil_law = fit_law(z,psi_e)
    g = conductance.fit_property_from_spline(g, soil_law, 'height', 'psi_e')

    return g