import dolfin as df
import numpy as np
import os
from . import *
from common.io import mpi_is_root
from common.bcs import Fixed, Charged
import pickle
__author__ = "Gaute Linga"



class Wall(df.SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary


def problem():
    info_cyan("Taylor-Green vortex flow with two-phase electrohydrodynamics.")

    solutes = [["c_p",  1, 3., 1., 2., -2.],
               ["c_m", -1, 4., 2., 1., -1.]]

    # Default parameters to be loaded unless starting from checkpoint.
    parameters = dict(
        solver="basic",
        folder="results_taylorgreen",
        restart_folder=False,
        enable_NS=True,
        enable_PF=True,
        enable_EC=True,
        save_intv=5,
        stats_intv=5,
        checkpoint_intv=50,
        tstep=0,
        dt=0.001,
        t_0=0.,
        T=.25,
        N=32,
        solutes=solutes,
        Lx=2.*np.pi,
        Ly=2.*np.pi,
        concentration_init=1.,
        concentration_init_dev=0.5,
        #
        density=[3.0, 1.0],
        viscosity=[3., 5.],
        permittivity=[3., 4.],
        EC_scheme="NL2",
        use_iterative_solvers=False,
        grav_const=0.,
        c_cutoff=0.1,
        p_lagrange=False,
        #
        surface_tension=.1,
        interface_thickness=1.0/np.sqrt(2.0),
        pf_mobility_coeff=1.0,
        pf_init=1.
    )
    return parameters


def constrained_domain(Lx, Ly, **namespace):
    return None  # PeriodicBoundary(Lx, Ly)


def mesh(Lx=2.*np.pi, Ly=2.*np.pi, N=16, **namespace):
    m = df.RectangleMesh(df.Point(0., 0.), df.Point(Lx, Ly), N, N)
    return m


def initialize(Lx, Ly,
               surface_tension, interface_thickness,
               pf_init, pf_mobility_coeff,
               solutes,
               density, viscosity, permittivity,
               restart_folder,
               field_to_subspace,
               field_to_subproblem,
               concentration_init,
               concentration_init_dev,
               enable_NS, enable_PF, enable_EC,
               **namespace):
    """ Create the initial state. """
    w_init_field = dict()

    #import matplotlib.pyplot as plt

    if not restart_folder:
        exprs = reference(t=0., **vars())
        for field, expr in exprs.items():
            if field_to_subproblem[field][1] == -1:
                S = field_to_subspace[field]
            else:
                S = field_to_subspace[field].collapse()
            w_init_field[field] = df.interpolate(
                expr, S)
            #fig = df.plot(w_init_field[field])
            #plt.title(field)
            #plt.colorbar(fig)
            #plt.show()

    return w_init_field


def create_bcs(t_0, Lx, Ly,
               concentration_init, concentration_init_dev,
               surface_tension, interface_thickness,
               pf_init, pf_mobility_coeff,
               solutes, density, permittivity, viscosity,
               p_lagrange,
               enable_NS, enable_PF, enable_EC,
               **namespace):
    """ The boundaries and boundary conditions are defined here. """
    boundaries = dict(
        wall=[Wall()]
    )
    bcs = dict()
    for boundary_name in boundaries.keys():
        bcs[boundary_name] = dict()

    # Apply pointwise BCs e.g. to pin pressure.
    bcs_pointwise = dict()

    exprs = reference(t=t_0, **vars())

    for field, expr in exprs.items():
        if field != "p":
            bcs["wall"][field] = Fixed(expr)

    if enable_NS and not p_lagrange:
        bcs_pointwise["p"] = (
            0., ("x[0] < DOLFIN_EPS && x[1] < DOLFIN_EPS"))

    return boundaries, bcs, bcs_pointwise


def tstep_hook(t, tstep, stats_intv, statsfile, field_to_subspace,
               field_to_subproblem, subproblems, w_, **namespace):
    info_blue("Timestep = {}".format(tstep))


def start_hook(newfolder, **namespace):
    statsfile = os.path.join(newfolder, "Statistics/stats.dat")
    return dict(statsfile=statsfile)


def rhs_source(t, solutes,
               surface_tension, interface_thickness, pf_mobility_coeff, pf_init,
               density, viscosity, permittivity,
               concentration_init, concentration_init_dev,
               enable_PF, enable_NS, enable_EC,
               **namespace):
    """ Source term on the right hand side of the conservation equations. """
    with open("data/taylorgreen_rhs.dat", "rb") as infile:
        code = pickle.load(infile)
    
    replace_vals = replace_dict(concentration_init, concentration_init_dev,
                                surface_tension, interface_thickness,
                                density, permittivity, viscosity, solutes,
                                pf_mobility_coeff, pf_init,
                                enable_NS, enable_PF, enable_EC)

    keys = []
    q = dict()
    if enable_NS and "u" in code:
        keys.append("u")
    if enable_PF and "phi" in code:
        keys.append("phi")
    if enable_EC:
        for solute in solutes:
            if solute[0] in code:
                keys.append(solute[0])
        if "V" in code:
            keys.append("V")

    for key in keys:
        q[key] = df.Expression(code[key], t=t, degree=2, **replace_vals)

    return q


def reference(t,
              concentration_init, concentration_init_dev,
              surface_tension, interface_thickness,
              pf_init, pf_mobility_coeff,
              solutes, density, permittivity, viscosity,
              enable_NS, enable_PF, enable_EC,
              **namespace):
    """ This contains the analytical reference for convergence analysis. """
    with open("data/taylorgreen_reference.dat", "rb") as infile:
        code = pickle.load(infile)

    replace_vals = replace_dict(concentration_init, concentration_init_dev,
                                surface_tension, interface_thickness,
                                density, permittivity, viscosity, solutes,
                                pf_mobility_coeff, pf_init,
                                enable_NS, enable_PF, enable_EC)

    code_strings = dict()
    if enable_NS:
        code_strings["u"] = code["u"]
        code_strings["p"] = code["p"]
    if enable_PF:
        code_strings["phi"] = code["phi"]
        code_strings["g"] = code["g"]
    if enable_EC:
        for solute in solutes:
            code_strings[solute[0]] = code[solute[0]]
        code_strings["V"] = code["V"]

    expr = dict()
    for key, code_string in code_strings.items():
        expr[key] = df.Expression(code_string, t=t, degree=5,
                                  **replace_vals)
    return expr


def mean(a):
    """ Same as np.mean """
    return (a[0]+a[1])/2


def dmean(a):
    """ Consider moving somewhere else. """
    return (a[0]-a[1])/2


def replace_dict(concentration_init, concentration_init_dev,
                 surface_tension, interface_thickness,
                 density, permittivity, viscosity, solutes,
                 pf_mobility_coeff, pf_init,
                 enable_NS, enable_PF, enable_EC):
    diffusivity_p = solutes[0][2:4]
    diffusivity_m = solutes[1][2:4]
    solubility_p = solutes[0][4:6]
    solubility_m = solutes[1][4:6]
    replace_vals = dict(
        U0=1. if enable_NS else 0,
        c0=concentration_init if enable_EC else 0,
        chi=concentration_init_dev if enable_EC else 0,
        sigma_tilde=surface_tension*3./(2.*np.sqrt(2.)) if enable_PF else 0,
        eps=interface_thickness if enable_PF else 1.0/np.sqrt(2),
        veps=mean(permittivity) if enable_PF and enable_EC else permittivity[0],
        dveps=dmean(permittivity) if enable_PF and enable_EC else 0,
        D_p=mean(diffusivity_p) if enable_PF else diffusivity_p[0],
        dD_p=dmean(diffusivity_p) if enable_PF else 0,
        D_m=mean(diffusivity_m) if enable_PF else diffusivity_m[0],
        dD_m=dmean(diffusivity_m) if enable_PF else 0,
        beta_p=mean(solubility_p) if enable_PF else 0,  # solubility_p[0],
        dbeta_p=dmean(solubility_p) if enable_PF else 0,
        beta_m=mean(solubility_m) if enable_PF else 0,  # solubility_m[0],
        dbeta_m=dmean(solubility_m) if enable_PF else 0,
        mu=mean(viscosity) if enable_PF else viscosity[0],
        dmu=dmean(viscosity) if enable_PF else 0,
        rho=mean(density) if enable_PF else density[0],
        drho=dmean(density) if enable_PF else 0,
        M=pf_mobility_coeff if enable_PF else 0,
        Phi0=pf_init if enable_PF else 0
    )
    return replace_vals


def end_hook(x_, t,
             field_to_subspace,
             concentration_init, concentration_init_dev,
             surface_tension, interface_thickness,
             pf_init, pf_mobility_coeff,
             solutes, density, permittivity, viscosity,
             enable_NS, enable_PF, enable_EC,
             testing,
             **namespace):
    if testing:
        exprs = reference(**vars())

        output = "Final error norms:"
        for field in ["u", "phi", "c_p", "c_m", "V"]:
            f = df.project(x_[field],
                           field_to_subspace[field].collapse())
            f_ref = exprs[field]
            output += " {} = {:e}".format(field, df.errornorm(f_ref, f, "L2"))
        info(output)
