# (c) 2017 Gregor Mitscha-Baude
"PNPS solvers and visualization for arbitrary pore"

import numpy as np
import dolfin
import nanopores as nano
import nanopores.physics.simplepnps as simplepnps
import nanopores.tools.solvers as solvers
import nanopores.tools.fields as fields
from nanopores.geometries.allpores import get_geo

default = nano.Params(
geop = nano.Params(
    geoname = "wei",
    dim = 2,
    x0 = None,
    rMolecule = 0.5,
    receptor = None,
    lcMolecule = 0.1,
    lcCenter = 0.4,
),
physp = nano.Params(
    Qmol = -1.,
    bulkcon = 1000.,
    dnaqsdamp = .5,
    bV = -0.1,
    rDPore = .9,
    Membraneqs = -0.0,
    ahemqs = None,
    ahemuniformqs = False,
    posDTarget = True,
),
solverp = nano.Params(
    h = 1.,
    frac = 0.2,
    Nmax = 2e4,
    imax = 30,
    tol = 1e-2,
    cheapest = False,
    stokesiter = False, #True
    diffusivity_data = None,
    reconstruct = False,
    fluid = "fluid",
    hybrid = True,
))
defaultp = default.geop | default.physp

class Setup(solvers.Setup):
    default = default

    def __init__(self, create_geo=True,
                 geop=None, physp=None, solverp=None, **params):
        self.init_params(params, geop=geop, physp=physp, solverp=solverp)
        self.init_geo(create_geo)
        self.init_phys()

    def init_geo(self, create_geo=True):
        h = self.solverp.h
        re = self.solverp.reconstruct
        self.geo = get_geo(h=h, reconstruct=re,
                           **self.geop) if create_geo else None

    def init_phys(self):
        cyl = self.geop.dim == 2
        if self.geo is None:
            self.physp["rTarget"] = self.geop["rMolecule"]*1e-9
        self.phys = nano.Physics("pore_mol", self.geo, cyl=cyl, **self.physp)

    def prerefine(self, visualize=False):
        return prerefine(self, visualize=visualize)
    
def get_active_params(params):
    setup = Setup(create_geo=False, **params)
    return setup.active_params

class Plotter(object):
    def __init__(self, setup=None, dim=3):
        if setup is not None:
            self.geo = setup.geo
            self.dim = setup.phys.dim
        else:
            self.dim = dim
        if self.dim == 3:
            if setup is not None:
                R = self.geo.params["R"]
                if "Hbot" in self.geo.params:
                    Htop = self.geo.params["Htop"]
                    Hbot = self.geo.params["Hbot"]
                else:
                    H = self.geo.params["H"]
                    Htop = Hbot = H/2.
            else:
                R, Htop, Hbot = 10., 15., 15.
            self.mesh2D = nano.RectangleMesh([-R,-Hbot], [R, Htop],
                                             int(4*R), int(2.*(Htop + Hbot)))