"""
mfusb module.

A module for a rough estimate of solute mass balance within unsaturated zone,
making use of UZF results

"""

import sys
import os
import subprocess as sub
import numpy as np

from flopy.modflow import *


class ModflowUsb():
    """
    MODFLOW Unsaturated Solute Balance Class.

    Parameters
    ----------
    model : FloPy model object
        The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
        this package will be linked.
    celldim : cell dimension
    cell_dict : dictionary of cells where solute source is present.
        keys are cellid
        values a list [row,col], where row and col are row and column of that cell, respectively
    mu_dict : dictionary of 1st-order degradation factor
        keys are are cellid
        values the degradation factor for cell corresponding to cellid
    cin_dict  : dictionary of concentration at ground surface
        keys are stress periods (0-based)
        values an array of shape = (nrow, ncol)


    """

    def __init__(self, model, uzf, celldim = None, cell_dict = None, cin_dict = None, mu_dict = None):

        # Inputs
        self.model = model
        self.uzf = uzf
        self.celldim = celldim
        self.cell_dict = cell_dict
        self.cin_dict = cin_dict
        self.mu_dict = mu_dict

        # Internal inputs
        nrow, ncol, nlay, nper = model.nrow_ncol_nlay_nper
        self.nrow = nrow
        self.ncol = ncol
        self.nlay = nlay
        self.nper = nper

        # Initialize
        inf_dict = {}
        rch_dict = {}
        z_dict = {}
        lay_dict = {}
        cout_dict = {}


        for kper in range(nper):
            inf_dict[kper] = np.zeros(shape = (nrow, ncol))
            rch_dict[kper] = np.zeros(shape = (nrow, ncol))
            z_dict[kper] = np.zeros(shape = (nrow, ncol))
            lay_dict[kper] = np.zeros(shape = (nrow, ncol))
            cout_dict[kper] = np.zeros(shape = (nrow, ncol))


        # Make it as self
        self.inf_dict = inf_dict
        self.rch_dict = rch_dict
        self.z_dict = z_dict
        self.lay_dict = lay_dict

        # This is the output
        self.cout_dict = cout_dict


    def write_uzf_gages(self):

        # retrieve UZF package of Modflow model
        ml = self.model
        uzf = self.uzf
        # retrieve DIS package of Modflow model, to get perlen and nstep
        # disinput = ml.name + '.dis' #
        disinput = ml.model_ws + '/' + ml.name + '.dis'

        dis = ModflowDis.load(disinput, ml)

        # make self this info
        self.Nsteps = dis.nstp
        self.SPlength = dis.perlen

        # Get ThetaS from UZF array thts and make it self

        self.thts = uzf.thts.get_value()

        # ---
        nuzgag = len(self.cell_dict)
            # write a gage file for each cell
        pathoutput = ml.model_ws + '/'
        row_col_iftunit_iuzopt = []
        icl = 1
        for g in self.cell_dict.keys():
            nr = self.cell_dict[g][0]
            nc = self.cell_dict[g][1]
            # define UNIT of file as 900 + cellid
            uncell = 900 + int(g)
            row_col_iftunit_iuzopt.append([nr, nc, uncell , 2  ])
            #
            ml.external_fnames.append(ml.name + '.uzf'+ str(icl))
            ml.external_units.append(uncell)
            icl += 1

        # Update UZF FloPy package
        uzf.nuzgag = nuzgag
        uzf.row_col_iftunit_iuzopt = row_col_iftunit_iuzopt



        # ... and run again Modflow Model
        ml.write_input()
        try :
            ml.run_model()
        except:
            # if it does not work, do the following ...
            namefile = ml.name + '.nam'
            #print 'NAM' , namefile
            #print 'exe',  ml.exename
            #print 'directory', ml.model_ws
            proc = sub.Popen([ml.exename, namefile], stdin = sub.PIPE, stdout = sub.PIPE, stderr = sub.PIPE, cwd= ml.model_ws)




    def solute_fluxout(self, dt, cin, cold, qin, qout, z, mu, th):
        """

        Solve the mass balance difference equation, implicit in time,
        along [0,z(t)] in space

        """
        if z <= 0:
            # No unsaturated zone
            cout = 0
        else:
            cout = ( (1-mu*dt)*cold + dt*qin*cin/(z*th) )/(1+dt*qout/(z*th)  )
        return cout

    def run_model(self, Nsteps , SPlength):
        """
        Run the mass balance calculation

        Returns
        -------
        None

        """

        ml = self.model

        uzf = ml.get_package('UZF')
        dis = ml.get_package('DIS')
        nper = self.nper


        # Get THTS from UZF
        thts = uzf.thts.get_value()

        # cell area
        area = self.celldim*self.celldim
            # retrieve data from UZF output files

        icl = 1
        for cl in self.cell_dict.keys():
            nr = self.cell_dict[cl][0]
            nc = self.cell_dict[cl][1]
            mu = self.mu_dict[cl]
            basename = os.path.join(ml.model_ws, ml.name + '.uzf')
            # basename = ml.model_ws + '\\' + ml.name  + '.uzf'
            filename = basename + str(icl)
            #print filename

            fl = open(filename, 'r')
            skip_rows = 2

            l = fl.readlines()
            irow = skip_rows+ Nsteps[0]
            mylist = [float(f) for f in l[irow].split()]
            lay = mylist[0]
            z = mylist[3]
            qin = mylist[11]
            qout = mylist[12]
            self.inf_dict[0][nr,nc] = qin/area
            self.rch_dict[0][nr,nc] = qout/area
            self.z_dict[0][nr,nc] = z

            for kper in range(1, nper):
                irow = irow + Nsteps[kper]
                mylist = [float(f) for f in l[irow].split()]
                lay = mylist[0]
                z = mylist[3]
                qin = mylist[11]
                qout = mylist[12]
                self.inf_dict[kper][nr,nc] = qin/area
                self.rch_dict[kper][nr,nc] = qout/area
                self.z_dict[kper][nr,nc] = z

            fl.close()
            # Update cell iterator
            icl += 1


            # Compute the solute mass balance
        #for cl in cell_dict.keys():
            #nr = cell_dict[cl][0]
            #nc = cell_dict[cl][1]
            # First SP
            cin = self.cin_dict[0][nr,nc]
            qin = self.inf_dict[0][nr,nc]
            qout = self.rch_dict[0][nr,nc]
            z = self.z_dict[0][nr,nc]
            th = float(thts[nr,nc])
            cold = 0
            # run
            cout = self.solute_fluxout(SPlength[0], cin, cold, qin, qout, z, mu, th)
            # Save the output
            self.cout_dict[0][nr,nc] = cout

            for kper in range(1, nper):
                cin = self.cin_dict[kper][nr,nc]
                qin = self.inf_dict[kper][nr,nc]
                qout = self.rch_dict[kper][nr,nc]
                cold = self.cout_dict[kper-1][nr,nc]
                z = self.z_dict[kper][nr,nc]
                # run
                cout = self.solute_fluxout(SPlength[kper], cin, cold, qin, qout, z, mu, th)
                # Save the output
                self.cout_dict[kper][nr,nc] = cout
