from __future__ import print_function

import os
import re
import concurrent.futures

import numpy as np
import netCDF4
from mdtraj.geometry import _geometry
from mdtraj.geometry.sasa import _ATOMIC_RADII

try:
    from bpmfwfft import IO
    try:
        from bpmfwfft.util import c_is_in_grid, cdistance, c_containing_cube
        from bpmfwfft.util import c_cal_charge_grid_new
        from bpmfwfft.util import c_cal_potential_grid
        from bpmfwfft.util import c_cal_lig_sasa_grid
        from bpmfwfft.util import c_cal_lig_sasa_grids
    except:
        from util import c_is_in_grid, cdistance, c_containing_cube
        from util import c_cal_charge_grid_new
        from util import c_cal_potential_grid
        from util import c_cal_lig_sasa_grid
        from util import c_cal_lig_sasa_grids

except:
    import IO
    from util import c_is_in_grid, cdistance, c_containing_cube
    from util import c_cal_charge_grid_new
    from util import c_cal_potential_grid
    from util import c_cal_lig_sasa_grid
    from util import c_cal_lig_sasa_grids

def process_potential_grid_function(
        name,
        crd,
        origin_crd,
        grid_spacing,
        grid_counts,
        charges,
        prmtop_ljsigma,
        molecule_sasa,
        rho,
        sasa_core_scaling,
        sasa_surface_scaling,
        sasa_grid
):
    """
    gets called by cal_potential_grid and assigned to a new python process
    use cython to calculate electrostatic, LJa, LJr, SASAr, and SASAi grids
    and save them to nc file
    """
    print("calculating Receptor %s grid" % name)
    grid_x = np.linspace(
        origin_crd[0],
        origin_crd[0] + ((grid_counts[0]-1) * grid_spacing[0]),
        num=grid_counts[0]
    )
    grid_y = np.linspace(
        origin_crd[1],
        origin_crd[1] + ((grid_counts[1] - 1) * grid_spacing[1]),
        num=grid_counts[1]
    )
    grid_z = np.linspace(
        origin_crd[2],
        origin_crd[2] + ((grid_counts[2] - 1) * grid_spacing[2]),
        num=grid_counts[2]
    )
    uper_most_corner_crd = origin_crd + (grid_counts - 1.) * grid_spacing
    uper_most_corner = (grid_counts - 1)

    grid = c_cal_potential_grid(name, crd,
                                grid_x, grid_y, grid_z,
                                origin_crd, uper_most_corner_crd, uper_most_corner,
                                grid_spacing, grid_counts,
                                charges, prmtop_ljsigma, molecule_sasa, rho,
                                sasa_core_scaling, sasa_surface_scaling, sasa_grid)
    return grid

def process_charge_grid_function(
        name,
        crd,
        origin_crd,
        grid_spacing,
        eight_corner_shifts,
        six_corner_shifts,
        grid_counts,
        charges,
        prmtop_ljsigma,
        molecule_sasa,
        sasa_grid
):
    """
    gets called by cal_potential_grid and assigned to a new python process
    use cython to calculate electrostatic, LJa, LJr, SASAr, and SASAi grids
    and save them to nc file
    """
    print("calculating Ligand %s grid" % name)
    grid_x = np.linspace(
        origin_crd[0],
        origin_crd[0] + ((grid_counts[0]-1) * grid_spacing[0]),
        num=grid_counts[0]
    )
    grid_y = np.linspace(
        origin_crd[1],
        origin_crd[1] + ((grid_counts[1] - 1) * grid_spacing[1]),
        num=grid_counts[1]
    )
    grid_z = np.linspace(
        origin_crd[2],
        origin_crd[2] + ((grid_counts[2] - 1) * grid_spacing[2]),
        num=grid_counts[2]
    )
    uper_most_corner_crd = origin_crd + (grid_counts - 1.) * grid_spacing
    uper_most_corner = (grid_counts - 1)

    grid = c_cal_charge_grid_new(name, crd,
                                grid_x, grid_y, grid_z,
                                origin_crd, uper_most_corner_crd, uper_most_corner,
                                grid_spacing, eight_corner_shifts, six_corner_shifts,
                                grid_counts, charges, prmtop_ljsigma, molecule_sasa, sasa_grid)

    return grid

def is_nc_grid_good(nc_grid_file):
    """
    :param nc_grid_file: name of nc file
    :return: bool
    """
    if not os.path.exists(nc_grid_file):
        return False

    if os.path.getsize(nc_grid_file) == 0:
        return False

    nc_handle = netCDF4.Dataset(nc_grid_file, "r")
    nc_keys = nc_handle.variables.keys()
    grid_keys = Grid().get_allowed_keys()
    for key in grid_keys:
        if key not in nc_keys:
            return False
    return True

class Grid(object):
    """
    an abstract class that defines some common methods and data attributes
    working implementations are in LigGrid and RecGrid below
    """
    def __init__(self):
        self._grid = {}
        # self._grid_func_names   = ("SASAi", "electrostatic", "LJr", "LJa", "SASAr")  # calculate all grids
        # self._grid_func_names = ("SASAi", "SASAr")  # uncomment to only calculate SASA grids
        self._grid_func_names = ()  # don't calculate any grids, but make grid objects for testing
        cartesian_axes  = ("x", "y", "z")
        box_dim_names   = ("d0", "d1", "d2")
        others          = ("spacing", "counts", "origin", "lj_sigma_scaling_factor", "sasa_core_scaling",
                           "sasa_surface_scaling")
        self._grid_allowed_keys = self._grid_func_names + cartesian_axes + box_dim_names + others

        self._eight_corner_shifts = [np.array([i,j,k], dtype=int) for i in range(2) for j in range(2) for k in range(2)]
        self._eight_corner_shifts = np.array(self._eight_corner_shifts, dtype=int)

        self._six_corner_shifts = self._get_six_corner_shifts()

        self._nearest_neighbor_shifts = self._get_nearest_neighbor_shifts()

    def _get_six_corner_shifts(self):
        six_corner_shifts = []
        for i in [-1, 1]:
            six_corner_shifts.append(np.array([i,0,0], dtype=int))
            six_corner_shifts.append(np.array([0,i,0], dtype=int))
            six_corner_shifts.append(np.array([0,0,i], dtype=int))
        return np.array(six_corner_shifts, dtype=int)

    def _get_nearest_neighbor_shifts(self):
        nearest_neighbor_shifts = []
        for i in [-1, 1]:
            nearest_neighbor_shifts.append(np.array([i, 0, 0], dtype=int))
            nearest_neighbor_shifts.append(np.array([0, i, 0], dtype=int))
            nearest_neighbor_shifts.append(np.array([0, 0, i], dtype=int))
            nearest_neighbor_shifts.append(np.array([0, i, i], dtype=int))
            nearest_neighbor_shifts.append(np.array([i, 0, i], dtype=int))
            nearest_neighbor_shifts.append(np.array([i, i, 0], dtype=int))
        return np.array(nearest_neighbor_shifts, dtype=int)

    def _set_grid_key_value(self, key, value):
        """
        key:    str
        value:  any object
        """
        assert key in self._grid_allowed_keys, key + " is not an allowed key"

        if key not in self._grid_func_names:
            print(value)
        self._grid[key] = value
        return None

    def _load_prmtop(self, prmtop_file_name, lj_sigma_scaling_factor):
        """
        :param prmtop_file_name: str, name of AMBER prmtop file
        :param lj_sigma_scaling_factor: float, must have value in [0.5, 1.0].
        It is stored in self._grid["lj_sigma_scaling_factor"] as
        a array of shape (1,) for reason of saving to nc file.
         Experience says that 0.8 is good for protein-ligand calculations.
        :return: None
        """
        assert 0.5 <= lj_sigma_scaling_factor <= 1.0, "lj_sigma_scaling_factor is out of allowed range"
        self._prmtop = IO.PrmtopLoad(prmtop_file_name).get_parm_for_grid_calculation()
        self._prmtop["LJ_SIGMA"] *= lj_sigma_scaling_factor
        self._set_grid_key_value("lj_sigma_scaling_factor", np.array([lj_sigma_scaling_factor], dtype=float))
        return None

    def _load_inpcrd(self, inpcrd_file_name):
        self._crd = IO.InpcrdLoad(inpcrd_file_name).get_coordinates()
        natoms = self._prmtop["POINTERS"]["NATOM"]
        if (self._crd.shape[0] != natoms) or (self._crd.shape[1] != 3):
            raise RuntimeError("coordinates in %s has wrong shape"%inpcrd_file_name)
        return None

    def _move_molecule_to(self, location):
        """
        Move the center of mass of the molecule to location.
        location:   3-array.
        This method affects self._crd.
        """
        assert len(location) == 3, "location must have len 3"
        displacement = np.array(location, dtype=float) - self._get_molecule_center_of_mass()
        for atom_ind in range(len(self._crd)):
            self._crd[atom_ind] += displacement
        return None

    def _get_molecule_center_of_mass(self):
        """
        return the center of mass of self._crd
        """
        center_of_mass = np.zeros([3], dtype=float)
        masses = self._prmtop["MASS"]
        for atom_ind in range(len(self._crd)):
            center_of_mass += masses[atom_ind] * self._crd[atom_ind]
        total_mass = masses.sum()
        if total_mass == 0:
            raise RuntimeError("zero total mass")
        return center_of_mass / total_mass

    def _get_molecule_sasa(self, probe_radius, n_sphere_points):
        """
        return the per atom SASA of the target molecule
        """
        xyz = self._crd
        xyz = np.expand_dims(xyz, 0)
        # convert coordinates to nanometers for mdtraj
        xyz = xyz.astype(np.float32)/10.

        atom_radii = []
        for atom_label in self._prmtop["PDB_TEMPLATE"]["ATOM_NAME"]:
            try:
                atom_radii.append(_ATOMIC_RADII[str(atom_label).split("-", 0)[0][0]])
            except:
                atom_radii.append(_ATOMIC_RADII[str(atom_label).split("-", 0)[0:1][0].title()])
        radii = np.array(atom_radii, np.float32) + probe_radius
        dim1 = xyz.shape[1]
        atom_mapping = np.arange(dim1, dtype=np.int32)
        out = np.zeros((xyz.shape[0], dim1), dtype=np.float32)
        _geometry._sasa(xyz, radii, int(n_sphere_points), atom_mapping, out)

        return out

    def _get_corner_crd(self, corner):
        """
        corner: 3-array integers
        """
        i, j, k = corner
        return np.array([self._grid["x"][i], self._grid["y"][j], self._grid["z"][k]] , dtype=float)

    def _get_uper_most_corner(self):
        return np.array(self._grid["counts"] - 1, dtype=int)

    def _get_uper_most_corner_crd(self):
        uper_most_corner = self._get_uper_most_corner()
        return self._get_corner_crd(uper_most_corner)

    def _get_origin_crd(self):
        return self._get_corner_crd([0,0,0])

    def _initialize_convenient_para(self):
        self._origin_crd           = self._get_origin_crd()
        self._uper_most_corner_crd = self._get_uper_most_corner_crd()
        self._uper_most_corner     = self._get_uper_most_corner()
        self._spacing              = np.array([self._grid["d%d"%i][i] for i in range(3)], dtype=float)
        return None

    def _is_in_grid(self, atom_coordinate):
        """
        in grid means atom_coordinate >= origin_crd and atom_coordinate < uper_most_corner_crd
        :param atom_coordinate: 3-array of float
        :return: bool
        """
        return c_is_in_grid(atom_coordinate, self._origin_crd, self._uper_most_corner_crd)

    def _distance(self, corner, atom_coordinate):
        """
        corner: 3-array int
        atom_coordinate:    3-array of float
        return distance from corner to atom coordinate
        """
        corner_crd = self._get_corner_crd(corner)
        return cdistance(atom_coordinate, corner_crd)

    def _containing_cube(self, atom_coordinate):
        eight_corners, nearest_ind, furthest_ind = c_containing_cube(atom_coordinate, self._origin_crd,
                                                                     self._uper_most_corner_crd,
                                                                     self._spacing, self._eight_corner_shifts,
                                                                     self._grid["x"], self._grid["y"], self._grid["z"])
        return eight_corners, nearest_ind, furthest_ind

    def _is_row_in_matrix(self, row, matrix):
        for r in matrix:
            if (row == r).all():
                return True
        return False

    def get_grid_func_names(self):
        return self._grid_func_names

    def get_grids(self):
        return self._grid

    def get_crd(self):
        return self._crd

    def get_prmtop(self):
        return self._prmtop

    def get_charges(self):
        charges = dict()
        for key in ["CHARGE_E_UNIT", "R_LJ_CHARGE", "A_LJ_CHARGE"]:
            charges[key] = self._prmtop[key]
        return charges

    def get_natoms(self):
        return self._prmtop["POINTERS"]["NATOM"]

    def get_allowed_keys(self):
        return self._grid_allowed_keys

class LigGrid(Grid):
    """
    Calculate the "charge" part of the interaction energy.
    """
    def __init__(self, prmtop_file_name, lj_sigma_scaling_factor,
                       inpcrd_file_name, receptor_grid):
        """
        :param prmtop_file_name: str, name of AMBER prmtop file
        :param lj_sigma_scaling_factor: float
        :param inpcrd_file_name: str, name of AMBER coordinate file
        :param receptor_grid: an instance of RecGrid class.
        """
        Grid.__init__(self)
        grid_data = receptor_grid.get_grids()
        if grid_data["lj_sigma_scaling_factor"][0] != lj_sigma_scaling_factor:
            raise RuntimeError("lj_sigma_scaling_factor is %f but in receptor_grid, it is %f" %(
                                lj_sigma_scaling_factor, grid_data["lj_sigma_scaling_factor"][0]))

        entries = [key for key in grid_data.keys() if key not in self._grid_func_names]
        print("Copy entries from receptor_grid", entries)
        for key in entries:
            self._set_grid_key_value(key, grid_data[key])
        self._initialize_convenient_para()

        self._rec_FFTs = receptor_grid.get_FFTs()

        self._load_prmtop(prmtop_file_name, lj_sigma_scaling_factor)
        self._load_inpcrd(inpcrd_file_name)
        self._move_ligand_to_lower_corner()
        self._molecule_sasa = self._get_molecule_sasa(0.14, 960)

    def _move_ligand_to_lower_corner(self):
        """
        move ligand to near the grid lower corner
        store self._max_grid_indices and self._initial_com
        """
        spacing = self._grid["spacing"]
        lower_ligand_corner = np.array([self._crd[:,i].min() for i in range(3)], dtype=float) - 2.5*spacing
        lower_ligand_corner_grid_aligned = lower_ligand_corner - (spacing + lower_ligand_corner % spacing) #new grid aligned variable
        upper_ligand_corner = np.array([self._crd[:,i].max() for i in range(3)], dtype=float) + 2.5*spacing
        upper_ligand_corner_grid_aligned = upper_ligand_corner + (spacing - upper_ligand_corner % spacing) #new grid aligned variable
        #print("lower ligand corner grid aligned=", lower_ligand_corner_grid_aligned)
        #print("upper ligand corner grid aligned=", upper_ligand_corner_grid_aligned)
        #
        ligand_box_lengths = upper_ligand_corner_grid_aligned - lower_ligand_corner_grid_aligned
#        ligand_box_lengths = upper_ligand_corner - lower_ligand_corner
        #print("ligand_box_lengths=", ligand_box_lengths)
        if np.any(ligand_box_lengths < 0):
            raise RuntimeError("One of the ligand box lengths are negative")

        max_grid_indices = np.ceil(ligand_box_lengths / spacing)
        self._max_grid_indices = self._grid["counts"] - np.array(max_grid_indices, dtype=int)
        if np.any(self._max_grid_indices <= 1):
            raise RuntimeError("At least one of the max grid indices is <= one")

        #displacement = self._origin_crd - lower_ligand_corner
        displacement = self._origin_crd - lower_ligand_corner_grid_aligned #formerly lower_ligand_corner
        for atom_ind in range(len(self._crd)):
            self._crd[atom_ind] += displacement
        print(f"Ligand translated by {displacement}")
        self._displacement = displacement
        lower_corner_origin = np.array([self._crd[:,i].min() for i in range(3)], dtype=float) - 1.5*spacing
        print(lower_corner_origin)
        self._initial_com = self._get_molecule_center_of_mass()
        return None

    def _get_charges(self, name):
        assert name in self._grid_func_names, "%s is not allowed"%name

        if name == "electrostatic":
            return np.array(self._prmtop["CHARGE_E_UNIT"], dtype=float)
        elif name == "LJa":
            return np.array(self._prmtop["A_LJ_CHARGE"], dtype=float)
        elif name == "LJr":
            return np.array(self._prmtop["R_LJ_CHARGE"], dtype=float)
        elif name == "SASAi":
            return np.array([0], dtype=float)
        elif name == "SASAr":
            return np.array([0], dtype=float)
        else:
            raise RuntimeError("%s is unknown"%name)

    def _cal_charge_grid(self, name, sasai_grid):
        # charges = self._get_charges(name)
        grid_counts = np.copy(self._grid["counts"])
        # grid = c_cal_charge_grid_new(name, self._crd, self._grid["x"], self._grid["y"], self._grid["z"],
        #                         self._origin_crd, self._uper_most_corner_crd, self._uper_most_corner,
        #                         self._grid["spacing"], self._eight_corner_shifts, self._six_corner_shifts,
        #                         grid_counts, charges, self._prmtop["LJ_SIGMA"],
        #                         self._molecule_sasa, sasai_grid
        #                         )
        task_divisor = 1
        with concurrent.futures.ProcessPoolExecutor() as executor:
            futures = {}
            sasa_grid = np.empty((0, 0, 0))
            for name in self._grid_func_names:
                futures_array = []
                for i in range(task_divisor):
                    counts = np.copy(self._grid["counts"])
                    counts_x = counts[0] // task_divisor
                    if i == task_divisor - 1:
                        counts_x += counts[0] % task_divisor
                    counts[0] = counts_x

                    grid_start_x = i * (self._grid["counts"][0] // task_divisor)
                    origin = np.copy(self._origin_crd)
                    origin[0] = grid_start_x * self._grid["spacing"][0]

                    if name != "SASAr":
                        dummy_grid = np.empty((1, 1, 1), dtype=np.float64)
                        futures_array.append(executor.submit(
                            process_charge_grid_function,
                            name,
                            self._crd,
                            origin,
                            self._grid["spacing"],
                            self._eight_corner_shifts,
                            self._six_corner_shifts,
                            counts,
                            self._get_charges(name),
                            self._prmtop["LJ_SIGMA"],
                            self._molecule_sasa,
                            dummy_grid
                        ))
                    else:
                        futures_array.append(executor.submit(
                            process_charge_grid_function,
                            name,
                            self._crd,
                            origin,
                            self._grid["spacing"],
                            self._eight_corner_shifts,
                            self._six_corner_shifts,
                            counts,
                            self._get_charges(name),
                            self._prmtop["LJ_SIGMA"],
                            self._molecule_sasa,
                            sasa_grid
                        ))
                futures[name] = futures_array
                # if name == "SASAi":
                #     sasa_array = []
                #     for i in range(task_divisor):
                #         print(futures[name][i].result().shape)
                #         partial_sasa_grid = futures[name][i].result()
                #         sasa_array.append(partial_sasa_grid)
                #     sasa_grid = np.concatenate(tuple(sasa_array))
            for name in futures:
                grid_array = []
                for i in range(task_divisor):
                    returned_sasai_grid, partial_grid = futures[name][i].result()
                    grid_array.append(partial_grid)
                    sasa_grid = returned_sasai_grid
                grid = np.concatenate(tuple(grid_array), axis=0)
                if name == "SASAi":
                    sasa_grid = np.copy(grid)
                # self._write_to_nc(nc_handle, name, grid)
                # self._set_grid_key_value(name, grid)
                # self._set_grid_key_value(name, None)     # to save memory

        return grid

    def _cal_corr_func(self, grid_name):
        """
        :param grid_name: str
        :return: fft correlation function
        """
        assert grid_name in self._grid_func_names, "%s is not an allowed grid name"%grid_name

        dummy_grid = np.empty((1, 1, 1), dtype=np.float64)
        grid = self._cal_charge_grid(grid_name, dummy_grid)

        self._set_grid_key_value(grid_name, grid)
        corr_func = np.fft.fftn(self._grid[grid_name])
        self._set_grid_key_value(grid_name, None)           # to save memory

        corr_func = corr_func.conjugate()
        corr_func = np.fft.ifftn(self._rec_FFTs[grid_name] * corr_func)
        corr_func = np.real(corr_func)
        return corr_func

    def _cal_shape_complementarity(self):
        """
        :param grid_name: str
        :return: fft correlation function
        """
        print("Calculating shape complementarity.")
        dummy_grid = np.empty((1, 1, 1), dtype=np.float64)
        counts = self._grid["counts"]

        lig_sasai_grid = self._cal_charge_grid("SASAi", dummy_grid)
        lig_sasar_grid = self._cal_charge_grid("SASAr", lig_sasai_grid)
        lig_sasa_grid = np.add(lig_sasar_grid, lig_sasai_grid*1.j)

        # self._set_grid_key_value(grid_name, lig_sasa_grid)
        corr_func = np.fft.fftn(lig_sasa_grid)
        # self._set_grid_key_value(grid_name, None)           # to save memory

        rec_sasa_grid = self._rec_FFTs["SASA"]

        rec_sasa_fft = np.fft.fftn(rec_sasa_grid)

        corr_func = np.fft.ifftn(rec_sasa_fft * corr_func) * (1/(np.prod(counts)))
        corr_func = np.real(corr_func) - np.imag(corr_func)

        return corr_func

    def _do_forward_fft(self, grid_name):
        assert grid_name in self._grid_func_names, "%s is not an allowed grid name"%grid_name
        grid = self._cal_charge_grid(grid_name)
        self._set_grid_key_value(grid_name, grid)
        forward_fft = np.fft.fftn(self._grid[grid_name])
        self._set_grid_key_value(grid_name, None)           # to save memory
        return forward_fft

    def _cal_corr_funcs(self, grid_names):
        """
        :param grid_names: list of str
        :return:
        """
        assert type(grid_names) == list, "grid_names must be a list"

        grid_name = grid_names[0]
        forward_fft = self._do_forward_fft(grid_name)
        corr_func = self._rec_FFTs[grid_name] * forward_fft.conjugate()

        for grid_name in grid_names[1:]:
            forward_fft = self._do_forward_fft(grid_name)
            corr_func += self._rec_FFTs[grid_name] * forward_fft.conjugate()

        corr_func = np.fft.ifftn(corr_func)
        corr_func = np.real(corr_func)
        return corr_func

    def _cal_energies(self):
        """
        calculate interaction energies
        store self._meaningful_energies (1-array) and self._meaningful_corners (2-array)
        meaningful means no border-crossing and no clashing
        TODO
        """
        max_i, max_j, max_k = self._max_grid_indices
        # TODO figure out how to calculate new corr function using SASA grids
        # corr_func = self._cal_corr_func("SASAr")
        corr_func = self._cal_shape_complementarity()
        self._free_of_clash = (corr_func > 0)
        print("number of poses free of clash:", self._free_of_clash.shape)
        self._free_of_clash = self._free_of_clash[0:max_i, 0:max_j, 0:max_k]  # exclude positions where ligand crosses border
        print("Ligand positions excluding border crossers", self._free_of_clash.shape)
        self._meaningful_energies = np.zeros(self._grid["counts"], dtype=float)
        if np.any(self._free_of_clash):
            grid_names = [name for name in self._grid_func_names if name[:4] != "SASA"]
            for name in grid_names:
                self._meaningful_energies += self._cal_corr_func(name)
        # get crystal pose here, use i,j,k of crystal pose
        self._meaningful_energies = self._meaningful_energies[0:max_i, 0:max_j, 0:max_k] # exclude positions where ligand crosses border

        self._meaningful_energies = self._meaningful_energies[self._free_of_clash]         # exclude positions where ligand is in clash with receptor, become 1D array
        self._number_of_meaningful_energies = self._meaningful_energies.shape[0]

        return None

    def _cal_energies_NOT_USED(self):
        """
        calculate interaction energies
        store self._meaningful_energies (1-array) and self._meaningful_corners (2-array)
        meaningful means no boder-crossing and no clashing
        TODO
        """
        max_i, max_j, max_k = self._max_grid_indices

        corr_func = self._cal_corr_func("occupancy")
        self._free_of_clash = (corr_func  < 0.001)
        self._free_of_clash = self._free_of_clash[0:max_i, 0:max_j, 0:max_k]  # exclude positions where ligand crosses border

        if np.any(self._free_of_clash):
            grid_names = [name for name in self._grid_func_names if name != "occupancy"]
            self._meaningful_energies = self._cal_corr_funcs(grid_names)
        else:
            self._meaningful_energies = np.zeros(self._grid["counts"], dtype=float)

        self._meaningful_energies = self._meaningful_energies[0:max_i, 0:max_j, 0:max_k] # exclude positions where ligand crosses border
        self._meaningful_energies = self._meaningful_energies[self._free_of_clash]         # exclude positions where ligand is in clash with receptor, become 1D array
        self._number_of_meaningful_energies = self._meaningful_energies.shape[0]
        return None

    def _cal_meaningful_corners(self):
        """
        return grid corners corresponding to self._meaningful_energies
        """
        corners = np.where(self._free_of_clash)
        corners = np.array(corners, dtype=int)
        corners = corners.transpose()
        return corners

    def _place_ligand_crd_in_grid(self, molecular_coord):
        """
        molecular_coord:    2-array, new ligand coordinate
        """
        crd = np.array(molecular_coord, dtype=float)
        natoms = self._prmtop["POINTERS"]["NATOM"]
        if (crd.shape[0] != natoms) or (crd.shape[1] != 3):
            raise RuntimeError("Input coord does not have the correct shape.")
        self._crd = crd
        self._move_ligand_to_lower_corner()
        return None

    def cal_grids(self, molecular_coord=None):
        """
        molecular_coord:    2-array, new ligand coordinate
        compute charge grids, meaningful_energies, meaningful_corners for molecular_coord
        if molecular_coord==None, self._crd is used
        """
        if molecular_coord is not None:
            self._place_ligand_crd_in_grid(molecular_coord)
        else:
            self._move_ligand_to_lower_corner()         # this is just in case the self._crd is not at the right position

        self._cal_energies()
        return None

    def get_bpmf(self, kB=0.001987204134799235, temperature=300.0):
        """
        use self._meaningful_energies to calculate and return exponential mean
        """
        if len(self._meaningful_energies) == 0:
            return 0.

        beta = 1. / temperature / kB
        V_0 = 1661.

        nr_samples = self.get_number_translations()
        energies = -beta *  self._meaningful_energies
        e_max = energies.max()
        exp_mean = np.exp(energies - e_max).sum() / nr_samples

        bpmf = -temperature * kB * (np.log(exp_mean) + e_max)

        V_binding = self.get_box_volume()
        correction = -temperature * kB * np.log(V_binding / V_0 / 8 / np.pi**2)
        return bpmf + correction

    def get_number_translations(self):
        return self._max_grid_indices.prod()

    def get_box_volume(self):
        """
        in angstrom ** 3
        """
        spacing = self._grid["spacing"]
        volume = ((self._max_grid_indices - 1) * spacing).prod()
        return volume

    def get_meaningful_energies(self):
        return self._meaningful_energies

    def get_meaningful_corners(self):
        meaningful_corners = self._cal_meaningful_corners()
        if meaningful_corners.shape[0] != self._number_of_meaningful_energies:
            raise RuntimeError("meaningful_corners does not have the same len as self._number_of_meaningful_energies")
        return meaningful_corners

    def set_meaningful_energies_to_none(self):
        self._meaningful_energies = None
        return None

    def get_initial_com(self):
        return self._initial_com

    def get_SASAi_grid(self, name, crd,
                       grid_x, grid_y, grid_z,
                       origin_crd, uper_most_corner_crd, uper_most_corner,
                       grid_spacing, eight_corner_shifts, six_corner_shifts,
                       grid_counts, charges, prmtop_ljsigma, molecule_sasa, sasa_grid,
                       roh_i_corners):
        sasai_grid, grid, roh_i_corners = c_cal_lig_sasa_grid(name, crd,
                                     grid_x, grid_y, grid_z,
                                     origin_crd, uper_most_corner_crd, uper_most_corner,
                                     grid_spacing, eight_corner_shifts, six_corner_shifts,
                                     grid_counts, charges, prmtop_ljsigma, molecule_sasa, sasa_grid,
                                               roh_i_corners)
        return sasai_grid, grid, roh_i_corners

    def get_SASAr_grid(self, name, crd,
                       grid_x, grid_y, grid_z,
                       origin_crd, uper_most_corner_crd, uper_most_corner,
                       grid_spacing, eight_corner_shifts, six_corner_shifts,
                       grid_counts, charges, prmtop_ljsigma, molecule_sasa, sasa_grid,
                       roh_i_corners):
        sasai_grid, grid, roh_i_corners = c_cal_lig_sasa_grid(name, crd,
                                     grid_x, grid_y, grid_z,
                                     origin_crd, uper_most_corner_crd, uper_most_corner,
                                     grid_spacing, eight_corner_shifts, six_corner_shifts,
                                     grid_counts, charges, prmtop_ljsigma, molecule_sasa, sasa_grid,
                                                              roh_i_corners)
        return sasai_grid, grid, roh_i_corners

    def get_SASA_grids(self, name, crd,
                       grid_x, grid_y, grid_z,
                       origin_crd, uper_most_corner_crd, uper_most_corner,
                       grid_spacing, eight_corner_shifts, six_corner_shifts,
                       nearest_neighbor_shifts, grid_counts, charges,
                       prmtop_ljsigma, molecule_sasa, rho,
                       sasa_core_scaling, sasa_surface_scaling):
        """
        Return the SASAi and SASAr grids for the Ligand
        """
        sasai_grid, sasar_grid = c_cal_lig_sasa_grids(name, crd,
                                     grid_x, grid_y, grid_z,
                                     origin_crd, uper_most_corner_crd, uper_most_corner,
                                     grid_spacing, eight_corner_shifts, six_corner_shifts,
                                     nearest_neighbor_shifts, grid_counts, charges,
                                     prmtop_ljsigma, molecule_sasa, rho,
                                     sasa_core_scaling, sasa_surface_scaling)
        return sasai_grid, sasar_grid

    def translate_ligand(self, displacement):
        """
        translate the ligand by displacement in Angstroms
        """
        for atom_ind in range(len(self._crd)):
            self._crd[atom_ind] += displacement
        return None

class RecGrid(Grid):
    """
    calculate the potential part of the interaction energy.
    """
    def __init__(self,  prmtop_file_name, lj_sigma_scaling_factor,
                        sasa_core_scaling, sasa_surface_scaling,
                        rho,
                        inpcrd_file_name,
                        bsite_file,
                        grid_nc_file,
                        new_calculation=False,
                        spacing=0.25, extra_buffer=3.0):  #default extra_buffer=3.0
        """
        :param prmtop_file_name: str, name of AMBER prmtop file
        :param lj_sigma_scaling_factor: float
        :param inpcrd_file_name: str, name of AMBER coordinate file
        :param bsite_file: str or None, if not None, name of a file defining the box dimension.
        This file is the same as "measured_binding_site.py" from AlGDock pipeline.
        :param grid_nc_file: str, name of grid nc file
        :param new_calculation: bool, if True do the new grid calculation else load data in grid_nc_file.
        :param spacing: float and in angstrom.
        :param extra_buffer: float
        """
        Grid.__init__(self)
        self._load_prmtop(prmtop_file_name, lj_sigma_scaling_factor)
        self._FFTs = {}

        if new_calculation:
            self._load_inpcrd(inpcrd_file_name)
            self._molecule_sasa = self._get_molecule_sasa(0.14, 960)
            self._rho = rho
            self._sasa_core_scaling = sasa_core_scaling
            self._sasa_surface_scaling = sasa_surface_scaling
            nc_handle = netCDF4.Dataset(grid_nc_file, "w", format="NETCDF4")
            self._write_to_nc(nc_handle, "lj_sigma_scaling_factor",
                                np.array([lj_sigma_scaling_factor], dtype=float))
            self._write_to_nc(nc_handle, "sasa_core_scaling",
                              np.array([sasa_core_scaling], dtype=float))
            self._write_to_nc(nc_handle, "sasa_surface_scaling",
                              np.array([sasa_surface_scaling], dtype=float))
            self._write_to_nc(nc_handle, "rho",
                              np.array([rho], dtype=float))
            self._write_to_nc(nc_handle, "molecule_sasa",
                              np.array(self._molecule_sasa, dtype=float))

            if bsite_file is not None:
                print("Receptor is assumed to be correctly translated such that box encloses binding pocket.")
                self._cal_grid_parameters_with_bsite(spacing, bsite_file, nc_handle)
                self._cal_grid_coordinates(nc_handle)
                self._initialize_convenient_para()
            else:
                print("No binding site specified, box encloses the whole receptor")
                self._cal_grid_parameters_without_bsite(spacing, extra_buffer, nc_handle)
                self._cal_grid_coordinates(nc_handle)
                self._initialize_convenient_para()
                self._move_receptor_to_grid_center()
                self._write_to_nc(nc_handle, "displacement", self._displacement)

            self._cal_potential_grids(nc_handle)
            self._write_to_nc(nc_handle, "trans_crd", self._crd)
            nc_handle.close()

        self._load_precomputed_grids(grid_nc_file, lj_sigma_scaling_factor)

    def _load_precomputed_grids(self, grid_nc_file, lj_sigma_scaling_factor):
        """
        nc_file_name:   str
        lj_sigma_scaling_factor: float, used for consistency check
        load netCDF file, populate self._grid with all the data fields
        """
        assert os.path.isfile(grid_nc_file), "%s does not exist" %grid_nc_file

        print(grid_nc_file)
        nc_handle = netCDF4.Dataset(grid_nc_file, "r")
        keys = [key for key in self._grid_allowed_keys if key not in self._grid_func_names]
        for key in keys:
            self._set_grid_key_value(key, nc_handle.variables[key][:])

        if self._grid["lj_sigma_scaling_factor"][0] != lj_sigma_scaling_factor:
            raise RuntimeError("lj_sigma_scaling_factor is %f but in %s, it is %f" %(
                lj_sigma_scaling_factor, grid_nc_file, self._grid["lj_sigma_scaling_factor"][0]))

        self._initialize_convenient_para()

        natoms = self._prmtop["POINTERS"]["NATOM"]
        if natoms != nc_handle.variables["trans_crd"].shape[0]:
            raise RuntimeError("Number of atoms is wrong in %s %nc_file_name")
        self._crd = nc_handle.variables["trans_crd"][:]

        for key in self._grid_func_names:
            if key[:4] != "SASA":
                self._set_grid_key_value(key, nc_handle.variables[key][:])
                self._FFTs[key] = self._cal_FFT(key)
                self._set_grid_key_value(key, None)     # to save memory
        # self._set_grid_key_value("SASAi", nc_handle.variables["SASAi"][:])  #UNCOMMENT ME
        # self._set_grid_key_value("SASAr", nc_handle.variables["SASAr"][:])  #UNCOMMENT ME
        # self._FFTs["SASA"] = self._cal_SASA_FFT()  #UNCOMMENT ME
        # self._set_grid_key_value("SASAi", None)  #UNCOMMENT ME
        # self._set_grid_key_value("SASAr", None)  #UNCOMMENT ME
        nc_handle.close()
        return None

    def _cal_FFT(self, name):
        if name not in self._grid_func_names:
            raise RuntimeError("%s is not allowed.")
        print("Doing FFT for %s"%name)
        FFT = np.fft.fftn(self._grid[name])
        return FFT

    def _cal_SASA_FFT(self):
        print("Doing FFT for SASA")
        sasai_grid = self._grid["SASAi"]
        sasar_grid = self._grid["SASAr"]
        sasa_grid = np.add(sasar_grid, sasai_grid*1.j)
        FFT = np.fft.fftn(sasa_grid)
        return FFT

    def _write_to_nc(self, nc_handle, key, value):
        print("Writing %s into nc file"%key)
        # create dimensions
        for dim in value.shape:
            dim_name = "%d"%dim
            if dim_name not in nc_handle.dimensions.keys():
                nc_handle.createDimension(dim_name, dim)

        # create variable
        if value.dtype == int:
            store_format = "i8"
        elif value.dtype == float:
            store_format = "f8"
        else:
            raise RuntimeError("unsupported dtype %s"%value.dtype)
        dimensions = tuple(["%d"%dim for dim in value.shape])
        nc_handle.createVariable(key, store_format, dimensions)

        # save data
        nc_handle.variables[key][:] = value
        return None

    def _cal_grid_parameters_with_bsite(self, spacing, bsite_file, nc_handle):
        """
        :param spacing: float, unit in angstrom, the same in x, y, z directions
        :param bsite_file: str, the file name of "measured_binding_site.py" from AlGDock pipeline
        :param nc_handle: an instance of netCDF4.Dataset()
        :return: None
        """
        assert spacing > 0, "spacing must be positive"
        self._set_grid_key_value("origin", np.zeros([3], dtype=float))

        self._set_grid_key_value("d0", np.array([spacing, 0, 0], dtype=float))
        self._set_grid_key_value("d1", np.array([0, spacing, 0], dtype=float))
        self._set_grid_key_value("d2", np.array([0, 0, spacing], dtype=float))
        self._set_grid_key_value("spacing", np.array([spacing]*3, dtype=float))

        # function to easily grab a single float from a complex string
        def get_num(x):
            return float(''.join(ele for ele in x if ele.isdigit() or ele == '.'))

        # create a regular expression to parse the read lines
        parser = re.compile(r'\d+.\d+')

        for line in open(bsite_file, "r"):
            if line.startswith('com_min = '):
                com_min = [float(i) for i in parser.findall(line)]
            if line.startswith('com_max = '):
                com_max = [float(i) for i in parser.findall(line)]
            if line.startswith('site_R = '):
                site_R = [float(i) for i in parser.findall(line)][0]
            if line.startswith('half_edge_length = '):
                half_edge_length = [float(i) for i in parser.findall(line)][0]
        #half_edge_length = get_num(line)
        print("half_edge_length = ", half_edge_length)
        length = 2. * half_edge_length         # TODO: this is not good, half_edge_length is define in bsite_file
        count = np.ceil(length / spacing) + 1

        self._set_grid_key_value("counts", np.array([count]*3, dtype=int))

        for key in ["origin", "d0", "d1", "d2", "spacing", "counts"]:
            self._write_to_nc(nc_handle, key, self._grid[key])
        return None

    def _cal_grid_parameters_without_bsite(self, spacing, extra_buffer, nc_handle):
        """
        use this when making box encompassing the whole receptor
        spacing:    float, unit in angstrom, the same in x, y, z directions
        extra_buffer: float
        """
        assert spacing > 0 and extra_buffer > 0, "spacing and extra_buffer must be positive"
        self._set_grid_key_value("origin", np.zeros( [3], dtype=float))

        self._set_grid_key_value("d0", np.array([spacing, 0, 0], dtype=float))
        self._set_grid_key_value("d1", np.array([0, spacing, 0], dtype=float))
        self._set_grid_key_value("d2", np.array([0, 0, spacing], dtype=float))
        self._set_grid_key_value("spacing", np.array([spacing]*3, dtype=float))

        lj_radius = np.array(self._prmtop["LJ_SIGMA"]/2., dtype=float)
        dx = (self._crd[:,0] + lj_radius).max() - (self._crd[:,0] - lj_radius).min()
        dy = (self._crd[:,1] + lj_radius).max() - (self._crd[:,1] - lj_radius).min()
        dz = (self._crd[:,2] + lj_radius).max() - (self._crd[:,2] - lj_radius).min()

        print("Receptor enclosing box [%f, %f, %f]"%(dx, dy, dz))
        print("extra_buffer: %f"%extra_buffer)

        length = max([dx, dy, dz]) + 2.0*extra_buffer

        if np.ceil(length / spacing)%2 != 0:
            length = length + spacing
        count = np.ceil(length / spacing) + 1

        self._set_grid_key_value("counts", np.array([count]*3, dtype=int))
        print("counts ", self._grid["counts"])
        print("Total box size %f" %((count-1)*spacing))

        for key in ["origin", "d0", "d1", "d2", "spacing", "counts"]:
            self._write_to_nc(nc_handle, key, self._grid[key])
        return None

    def _move_receptor_to_grid_center(self):
        """
        use this when making box encompassing the whole receptor
        """
        spacing = self._grid["spacing"]
        lower_receptor_corner = np.array([self._crd[:,i].min() for i in range(3)], dtype=float)
        upper_receptor_corner = np.array([self._crd[:,i].max() for i in range(3)], dtype=float)

        lower_receptor_corner_grid_aligned = lower_receptor_corner - (spacing + lower_receptor_corner % spacing)
        upper_receptor_corner_grid_aligned = upper_receptor_corner + (spacing - upper_receptor_corner % spacing)

        receptor_box_center_grid_aligned = (upper_receptor_corner_grid_aligned + lower_receptor_corner_grid_aligned) / 2.

        receptor_box_center = (upper_receptor_corner + lower_receptor_corner) / 2.

        total_grid_count = (self._uper_most_corner_crd+spacing)/spacing
        print(total_grid_count)
        grid_center = (self._origin_crd + self._uper_most_corner_crd) / 2.
        receptor_box_length = upper_receptor_corner - lower_receptor_corner
        receptor_box_length_grid_aligned = upper_receptor_corner_grid_aligned - lower_receptor_corner_grid_aligned

        #test redefs of variables
#        receptor_box_center = ([upper_receptor_corner_grid_aligned[0],
#            upper_receptor_corner_grid_aligned[1]+0.5,
#            upper_receptor_corner_grid_aligned[2]+0.5] + lower_receptor_corner_grid_aligned) / 2.
        for index, coord in enumerate(upper_receptor_corner_grid_aligned):
            corner_to_corner_1D_distance = (coord - lower_receptor_corner_grid_aligned[index])/spacing[index]
            lower_corner_coord = lower_receptor_corner_grid_aligned[index]
            half_spacing = spacing[index]/2.
            print(corner_to_corner_1D_distance)
            if corner_to_corner_1D_distance%2 == 0:
                shifted_upper_coord = coord + half_spacing
                shifted_lower_coord = lower_corner_coord - half_spacing
                upper_receptor_corner_grid_aligned[index] = shifted_upper_coord
                lower_receptor_corner_grid_aligned[index] = shifted_lower_coord

        receptor_box_center = (upper_receptor_corner_grid_aligned + lower_receptor_corner_grid_aligned) / 2.
        grid_snap = np.mod(receptor_box_center, spacing)
        if np.any(np.where(grid_snap != 0)):
            receptor_box_center = np.add(receptor_box_center, np.subtract(spacing, grid_snap))

        print('receptor_box_center', receptor_box_center)
        displacement = grid_center - receptor_box_center

        print('lower_receptor_corner_grid_aligned: ', lower_receptor_corner_grid_aligned,
            '\nupper_receptor_corner_grid_aligned: ', upper_receptor_corner_grid_aligned,
            '\nlower_receptor_corner: ', lower_receptor_corner,
            '\nupper_receptor_corner: ', upper_receptor_corner,
            '\nreceptor_box_center: ', receptor_box_center,
            '\nreceptor_box_center_grid_aligned', receptor_box_center_grid_aligned,
            '\ngrid_center: ', grid_center,
            '\nreceptor_box_length: ', receptor_box_length,
            '\nreceptor_box_length_grid_aligned: ', receptor_box_length_grid_aligned,
            '\nspacing num', receptor_box_length_grid_aligned/spacing
            )
        print("Receptor is translated by ", displacement)
        self._displacement = displacement

        for atom_ind in range(len(self._crd)):
            self._crd[atom_ind] += displacement
        return None

    def _cal_grid_coordinates(self, nc_handle):
        """
        calculate grid coordinates (x,y,z) for each corner,
        save 'x', 'y', 'z' to self._grid
        """
        print("calculating grid coordinates")
        #
        x = np.zeros(self._grid["counts"][0], dtype=float)
        y = np.zeros(self._grid["counts"][1], dtype=float)
        z = np.zeros(self._grid["counts"][2], dtype=float)

        for i in range(self._grid["counts"][0]):
            x[i] = self._grid["origin"][0] + i*self._grid["d0"][0]

        for j in range(self._grid["counts"][1]):
            y[j] = self._grid["origin"][1] + j*self._grid["d1"][1]

        for k in range(self._grid["counts"][2]):
            z[k] = self._grid["origin"][2] + k*self._grid["d2"][2]

        self._set_grid_key_value("x", x)
        self._set_grid_key_value("y", y)
        self._set_grid_key_value("z", z)

        for key in ["x", "y", "z"]:
            self._write_to_nc(nc_handle, key, self._grid[key])
        return None

    def _get_charges(self, name):
        assert name in self._grid_func_names, "%s is not allowed"%name

        if name == "electrostatic":
            return 332.05221729 * np.array(self._prmtop["CHARGE_E_UNIT"], dtype=float)
        elif name == "LJa":
            return -2.0 * np.array(self._prmtop["A_LJ_CHARGE"], dtype=float)
        elif name == "LJr":
            return np.array(self._prmtop["R_LJ_CHARGE"], dtype=float)
        elif name == "SASAi":
            return np.array([0], dtype=float)
        elif name == "SASAr":
            return np.array([0], dtype=float)
        else:
            raise RuntimeError("%s is unknown"%name)

    def _cal_potential_grids(self, nc_handle):
        """
        Divides each grid calculation into a separate process (electrostatic, LJr, LJa,
        SASAr, SASAi) and then divides the grid into slices along the x-axis determined by
        the "task divisor". Remainders are calculated in the last slice.  This adds
        multiprocessing functionality to the grid generation.
        """
        task_divisor = 8
        with concurrent.futures.ProcessPoolExecutor() as executor:
            futures = {}
            sasa_grid = np.empty((0,0,0))
            for name in self._grid_func_names:
                futures_array = []
                for i in range(task_divisor):
                    counts = np.copy(self._grid["counts"])
                    counts_x = counts[0] // task_divisor
                    if i == task_divisor-1:
                        counts_x += counts[0] % task_divisor
                    counts[0] = counts_x

                    grid_start_x = i * (self._grid["counts"][0] // task_divisor)
                    origin = np.copy(self._origin_crd)
                    origin[0] = grid_start_x * self._grid["spacing"][0]

                    if name != "SASAr":
                        dummy_grid = np.empty((1,1,1), dtype=np.float64)
                        futures_array.append(executor.submit(
                            process_potential_grid_function,
                            name,
                            self._crd,
                            origin,
                            self._grid["spacing"],
                            counts,
                            self._get_charges(name),
                            self._prmtop["LJ_SIGMA"],
                            self._molecule_sasa,
                            self._rho,
                            self._sasa_core_scaling,
                            self._sasa_surface_scaling,
                            dummy_grid
                        ))
                    else:
                        futures_array.append(executor.submit(
                            process_potential_grid_function,
                            name,
                            self._crd,
                            origin,
                            self._grid["spacing"],
                            counts,
                            self._get_charges(name),
                            self._prmtop["LJ_SIGMA"],
                            self._molecule_sasa,
                            self._rho,
                            self._sasa_core_scaling,
                            self._sasa_surface_scaling,
                            sasa_grid
                        ))
                futures[name] = futures_array
                if name == "SASAi":
                    sasa_array = []
                    for i in range(task_divisor):
                        partial_sasa_grid = futures[name][i].result()
                        sasa_array.append(partial_sasa_grid)
                    sasa_grid = np.concatenate(tuple(sasa_array))
            for name in futures:
                grid_array = []
                for i in range(task_divisor):
                    partial_grid = futures[name][i].result()
                    grid_array.append(partial_grid)
                grid = np.concatenate(tuple(grid_array), axis=0)
                if name == "SASAi":
                    sasa_grid = np.copy(grid)
                self._write_to_nc(nc_handle, name, grid)
                self._set_grid_key_value(name, grid)
                # self._set_grid_key_value(name, None)     # to save memory

        return None

    def _exact_values(self, coordinate):
        """
        coordinate: 3-array of float
        calculate the exact "potential" value at any coordinate
        """
        assert len(coordinate) == 3, "coordinate must have len 3"
        if not self._is_in_grid(coordinate):
            raise RuntimeError("atom is outside grid even after pbc translated")

        values = {}
        for name in self._grid_func_names:
            if name[:4] != "SASA":
                values[name] = 0.

        NATOM = self._prmtop["POINTERS"]["NATOM"]
        for atom_ind in range(NATOM):
            dif = coordinate - self._crd[atom_ind]
            R = np.sqrt((dif*dif).sum())
            lj_diameter = self._prmtop["LJ_SIGMA"][atom_ind]

            if R > lj_diameter:
                values["electrostatic"] +=  332.05221729 * self._prmtop["CHARGE_E_UNIT"][atom_ind] / R
                values["LJr"] +=  self._prmtop["R_LJ_CHARGE"][atom_ind] / R**12
                values["LJa"] += -2. * self._prmtop["A_LJ_CHARGE"][atom_ind] / R**6

        return values

    def _trilinear_interpolation( self, grid_name, coordinate ):
        """
        grid_name is a str one of "electrostatic", "LJr" and "LJa"
        coordinate is an array of three numbers
        trilinear interpolation
        https://en.wikipedia.org/wiki/Trilinear_interpolation
        """
        raise RuntimeError("Do not use, not tested yet")
        assert len(coordinate) == 3, "coordinate must have len 3"

        eight_corners, nearest_ind, furthest_ind = self._containing_cube( coordinate ) # throw exception if coordinate is outside
        lower_corner = eight_corners[0]

        (i0, j0, k0) = lower_corner
        (i1, j1, k1) = (i0 + 1, j0 + 1, k0 + 1)

        xd = (coordinate[0] - self._grid["x"][i0,j0,k0]) / (self._grid["x"][i1,j1,k1] - self._grid["x"][i0,j0,k0])
        yd = (coordinate[1] - self._grid["y"][i0,j0,k0]) / (self._grid["y"][i1,j1,k1] - self._grid["y"][i0,j0,k0])
        zd = (coordinate[2] - self._grid["z"][i0,j0,k0]) / (self._grid["z"][i1,j1,k1] - self._grid["z"][i0,j0,k0])

        c00 = self._grid[grid_name][i0,j0,k0]*(1. - xd) + self._grid[grid_name][i1,j0,k0]*xd
        c10 = self._grid[grid_name][i0,j1,k0]*(1. - xd) + self._grid[grid_name][i1,j1,k0]*xd
        c01 = self._grid[grid_name][i0,j0,k1]*(1. - xd) + self._grid[grid_name][i1,j0,k1]*xd
        c11 = self._grid[grid_name][i0,j1,k1]*(1. - xd) + self._grid[grid_name][i1,j1,k1]*xd

        c0 = c00*(1. - yd) + c10*yd
        c1 = c01*(1. - yd) + c11*yd

        c = c0*(1. - zd) + c1*zd
        return c

    def direct_energy(self, ligand_coordinate, ligand_charges):
        """
        :param ligand_coordinate: ndarray of shape (natoms, 3)
        :param ligand_charges: ndarray of shape (3,)
        :return: dic
        """
        assert len(ligand_coordinate) == len(ligand_charges["CHARGE_E_UNIT"]), "coord and charges must have the same len"
        energy = 0.
        for atom_ind in range(len(ligand_coordinate)):
            potentials = self._exact_values(ligand_coordinate[atom_ind])
            energy += potentials["electrostatic"]*ligand_charges["CHARGE_E_UNIT"][atom_ind]
            energy += potentials["LJr"]*ligand_charges["R_LJ_CHARGE"][atom_ind]
            energy += potentials["LJa"]*ligand_charges["A_LJ_CHARGE"][atom_ind]
        return energy

    def interpolated_energy(self, ligand_coordinate, ligand_charges):
        """
        ligand_coordinate:  array of shape (natoms, 3)
        ligand_charges: array of shape (3)
        assume that ligand_coordinate is inside grid
        """
        raise RuntimeError("Do not use, not tested yet")
        assert len(ligand_coordinate) == len(ligand_charges["CHARGE_E_UNIT"]), "coord and charges must have the same len"
        grid_names = [name for name in self._grid_func_names if name[:4] != "SASA"]
        energy = 0.
        potentials = {}
        for atom_ind in range(len(ligand_coordinate)):
            for name in grid_names:
                potentials[name] = self._trilinear_interpolation(name, ligand_coordinate[atom_ind])

            energy += potentials["electrostatic"]*ligand_charges["CHARGE_E_UNIT"][atom_ind]
            energy += potentials["LJr"]*ligand_charges["R_LJ_CHARGE"][atom_ind]
            energy += potentials["LJa"]*ligand_charges["A_LJ_CHARGE"][atom_ind]

        return energy

    def get_FFTs(self):
        return self._FFTs

    def write_box(self, file_name):
        IO.write_box(self, file_name)
        return None

    def write_pdb(self, file_name, mode):
        IO.write_pdb(self._prmtop, self._crd, file_name, mode)
        return None

if __name__ == "__main__":
    # do some test
    rec_prmtop_file = "../examples/amber/ubiquitin_ligase/receptor.prmtop"
    rec_inpcrd_file = "../examples/amber/ubiquitin_ligase/receptor.inpcrd"
    grid_nc_file = "../examples/grid/ubiquitin_ligase/grid.nc"
    lj_sigma_scaling_factor = 0.8
    # bsite_file = "../examples/amber/t4_lysozyme/measured_binding_site.py"
    bsite_file = None
    spacing = 0.5

    rec_grid = RecGrid(rec_prmtop_file, lj_sigma_scaling_factor, rec_inpcrd_file,
                        bsite_file,
                        grid_nc_file,
                        new_calculation=True,
                        spacing=spacing)
    print("get_grid_func_names", rec_grid.get_grid_func_names())
    print("get_grids", rec_grid.get_grids())
    print("get_crd", rec_grid.get_crd())
    print("get_prmtop", rec_grid.get_prmtop())
    print("get_prmtop", rec_grid.get_charges())
    print("get_natoms", rec_grid.get_natoms())
    print("get_natoms", rec_grid.get_allowed_keys())

    rec_grid.write_box("../examples/grid/ubiquitin_ligase/box.pdb")
    rec_grid.write_pdb("../examples/grid/ubiquitin_ligase/test.pdb", "w")

    lig_prmtop_file = "../examples/amber/ubiquitin/ligand.prmtop"
    lig_inpcrd_file = "../examples/amber/ubiquitin/ligand.inpcrd"
    lig_grid = LigGrid(lig_prmtop_file, lj_sigma_scaling_factor, lig_inpcrd_file, rec_grid)
    lig_grid.cal_grids()
    print("get_bpmf", lig_grid.get_bpmf())
    print("get_number_translations", lig_grid.get_number_translations())
    print("get_box_volume", lig_grid.get_box_volume())
    print("get_meaningful_energies", lig_grid.get_meaningful_energies())
    print("get_meaningful_corners", lig_grid.get_meaningful_corners())
    print("set_meaningful_energies_to_none", lig_grid.set_meaningful_energies_to_none())
    print("get_initial_com", lig_grid.get_initial_com())
    print("Receptor SASA", rec_grid._get_molecule_sasa(0.14, 960))
    print("Ligand SASA", lig_grid._get_molecule_sasa(0.14, 960))


# This file is part of pyTSEB for running different TSEB models
# Copyright 2016 Hector Nieto and contributors listed in the README.md file.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

'''
Created on Apr 6 2015
@author: Hector Nieto (hnieto@ias.csic.es)
Modified on Jan 27 2016
@author: Hector Nieto (hnieto@ias.csic.es)
DESCRIPTION
===========
This package contains the main routines inherent of Two Source Energy Balance `TSEB` models.
Additional functions needed in TSEB, such as computing of net radiation or estimating the
resistances to heat and momentum transport are imported.
* :doc:`netRadiation` for the estimation of net radiation and radiation partitioning.
* :doc:`ClumpingIndex` for the estimatio of canopy clumping index.
* :doc:`meteoUtils` for the estimation of meteorological variables.
* :doc:`resistances` for the estimation of the resistances to heat and momentum transport.
* :doc:`MOsimilarity` for the estimation of the Monin-Obukhov length and MOST-related variables.
PACKAGE CONTENTS
================
TSEB models
-----------
* :func:`TSEB_PT` Priestley-Taylor TSEB using a single observation of composite radiometric temperature.
Ancillary functions
-------------------
* :func:`calc_F_theta_campbell`. Gap fraction estimation.
* :func:`calc_G_time_diff`. Santanello & Friedl (2003) [Santanello2003]_ soil heat flux model.
* :func:`calc_G_ratio`. Soil heat flux as a fixed fraction of net radiation [Choudhury1987]_.
* :func:`calc_H_C_PT`. Priestley- Taylor Canopy sensible heat flux.
* :func:`calc_T_C_series.` Canopy temperature from canopy sensible heat flux and resistance in series.
* :func:`calc_T_S`. Soil temperature from form composite radiometric temperature.
'''

import numpy as np
#from FourSAIL import FourSAIL

from .meteo_utils import calc_rho,calc_c_p,calc_delta_vapor_pressure,calc_psicr,calc_lambda
from .resistances import calc_z_0H,calc_R_A,calc_R_x_Norman,calc_R_S_Kustas,calc_R_x_Choudhury,calc_R_S_Choudhury,calc_R_x_McNaughton,calc_R_S_McNaughton
from .MO_similarity import calc_u_star,calc_u_C_star,calc_u_Goudriaan,calc_L,calc_A_Goudriaan
from .net_radiation import calc_L_n_Kustas,calc_K_be_Campbell
from .clumping_index import calc_omega0_Kustas
#from cachetools import cached

#==============================================================================
# List of constants used in TSEB model and sub-routines
#==============================================================================
# Change threshold in  Monin-Obukhov lengh to stop the iterations
L_thres = 0.00001
# Change threshold in  friction velocity to stop the iterations
u_thres = 0.00001
# mimimun allowed friction velocity
u_friction_min = 0.01
# Maximum number of interations
ITERATIONS = 20
# kB coefficient
kB = 0.0
# Stephan Boltzmann constant (W m-2 K-4)
sb = 5.670373e-8

def TSEB_PT(
    Tr_K,
    vza,
    T_A_K,
    u,
    ea,
    p,
    Sn_C,
    Sn_S,
    L_dn,
    LAI,
    hc,
    emis_C,
    emis_S,
    z_0M,
    d_0,
    z_u,
    z_T,
    nullMask,
    leaf_width=0.1,
    z0_soil=0.01,
    alpha_PT=1.26,
    x_LAD=1,
    f_c=1.0,
    f_g=1.0,
    w_C=1.0,
    resistance_form=0,
    calcG_array=[
        [1],
        0.35],
        UseL=False):
    '''Priestley-Taylor TSEB
    Calculates the Priestley Taylor TSEB fluxes using a single observation of
    composite radiometric temperature and using resistances in series.
    Parameters
    ----------
    Tr_K : float
        Radiometric composite temperature (Kelvin).
    vza : float
        View Zenith Angle (degrees).
    T_A_K : float
        Air temperature (Kelvin).
    u : float
        Wind speed above the canopy (m s-1).
    ea : float
        Water vapour pressure above the canopy (mb).
    p : float
        Atmospheric pressure (mb), use 1013 mb by default.
    Sn_C : float
        Canopy net shortwave radiation (W m-2).
    Sn_S : float
        Soil net shortwave radiation (W m-2).
    L_dn : float
        Downwelling longwave radiation (W m-2).
    LAI : float
        Effective Leaf Area Index (m2 m-2).
    hc : float
        Canopy height (m).
    emis_C : float
        Leaf emissivity.
    emis_S : flaot
        Soil emissivity.
    z_0M : float
        Aerodynamic surface roughness length for momentum transfer (m).
    d_0 : float
        Zero-plane displacement height (m).
    z_u : float
        Height of measurement of windspeed (m).
    z_T : float
        Height of measurement of air temperature (m).
    leaf_width : float, optional
        average/effective leaf width (m).
    z0_soil : float, optional
        bare soil aerodynamic roughness length (m).
    alpha_PT : float, optional
        Priestley Taylor coeffient for canopy potential transpiration,
        use 1.26 by default.
    x_LAD : float, optional
        Campbell 1990 leaf inclination distribution function chi parameter.
    f_c : float, optional
        Fractional cover.
    f_g : float, optional
        Fraction of vegetation that is green.
    w_C : float, optional
        Canopy width to height ratio.
    resistance_form : int, optional
        Flag to determine which Resistances R_x, R_S model to use.
            * 0 [Default] Norman et al 1995 and Kustas et al 1999.
            * 1 : Choudhury and Monteith 1988.
            * 2 : McNaughton and Van der Hurk 1995.
    calcG_params : list[list,float or array], optional
        Method to calculate soil heat flux,parameters.
            * [[1],G_ratio]: default, estimate G as a ratio of Rn_S, default Gratio=0.35.
            * [[0],G_constant] : Use a constant G, usually use 0 to ignore the computation of G.
            * [[2,Amplitude,phase_shift,shape],time] : estimate G from Santanello and Friedl with G_param list of parameters (see :func:`~TSEB.calc_G_time_diff`).
    UseL : float or None, optional
        If included, its value will be used to force the Moning-Obukhov stability length.
    Returns
    -------
    flag : int
        Quality flag, see Appendix for description.
    T_S : float
        Soil temperature  (Kelvin).
    T_C : float
        Canopy temperature  (Kelvin).
    T_AC : float
        Air temperature at the canopy interface (Kelvin).
    L_nS : float
        Soil net longwave radiation (W m-2)
    L_nC : float
        Canopy net longwave radiation (W m-2)
    LE_C : float
        Canopy latent heat flux (W m-2).
    H_C : float
        Canopy sensible heat flux (W m-2).
    LE_S : float
        Soil latent heat flux (W m-2).
    H_S : float
        Soil sensible heat flux (W m-2).
    G : float
        Soil heat flux (W m-2).
    R_S : float
        Soil aerodynamic resistance to heat transport (s m-1).
    R_x : float
        Bulk canopy aerodynamic resistance to heat transport (s m-1).
    R_A : float
        Aerodynamic resistance to heat transport (s m-1).
    u_friction : float
        Friction velocity (m s-1).
    L : float
        Monin-Obuhkov length (m).
    n_iterations : int
        number of iterations until convergence of L.
    References
    ----------
    .. [Norman1995] J.M. Norman, W.P. Kustas, K.S. Humes, Source approach for estimating
        soil and vegetation energy fluxes in observations of directional radiometric
        surface temperature, Agricultural and Forest Meteorology, Volume 77, Issues 3-4,
        Pages 263-293,
        http://dx.doi.org/10.1016/0168-1923(95)02265-Y.
    .. [Kustas1999] William P Kustas, John M Norman, Evaluation of soil and vegetation heat
        flux predictions using a simple two-source model with radiometric temperatures for
        partial canopy cover, Agricultural and Forest Meteorology, Volume 94, Issue 1,
        Pages 13-29,
        http://dx.doi.org/10.1016/S0168-1923(99)00005-2.
    '''

    # Convert input float scalars to arrays and parameters size
    Tr_K = np.asarray(Tr_K)
#    (vza,
#     T_A_K,
#     u,
#     ea,
#     p,
#     Sn_C,
#     Sn_S,
#     L_dn,
#     LAI,
#     hc,
#     emis_C,
#     emis_S,
#     z_0M,
#     d_0,
#     z_u,
#     z_T,
#     nullMask,
#     leaf_width,
#     z0_soil,
#     alpha_PT,
#     x_LAD,
#     f_c,
#     f_g,
#     w_C,
#     calcG_array) = map(_check_default_parameter_size,
#                        [vza,
#                         T_A_K,
#                         u,
#                         ea,
#                         p,
#                         Sn_C,
#                         Sn_S,
#                         L_dn,
#                         LAI,
#                         hc,
#                         emis_C,
#                         emis_S,
#                         z_0M,
#                         d_0,
#                         z_u,
#                         z_T,
#                         nullMask,
#                         leaf_width,
#                         z0_soil,
#                         alpha_PT,
#                         x_LAD,
#                         f_c,
#                         f_g,
#                         w_C,
#                         calcG_params[1]],
#                        [Tr_K] * 24)
    calcG_array[1] = _check_default_parameter_size(calcG_array[1],Tr_K)
    # Create the output variables
    [flag, T_S, T_C, T_AC, Ln_S, Ln_C, LE_C, H_C, LE_S, H_S, G, R_S, R_x,
        R_A, iterations] = [np.zeros(Tr_K.shape) for i in range(15)]

    # iteration of the Monin-Obukhov length
    if isinstance(UseL, bool):
        # Initially assume stable atmospheric conditions and set variables for
        L = np.asarray(np.zeros(T_S.shape) + np.inf)
        L = np.asarray(np.zeros(T_S.shape))
        max_iterations = ITERATIONS
    else:  # We force Monin-Obukhov lenght to the provided array/value
        L = np.asarray(np.ones(T_S.shape) * UseL)
        max_iterations = 1  # No iteration
    # Calculate the general parameters
    if ea.sum() == 0.0:
        z = 350.0
        rho = 101.3*((((T_A_K)-(0.0065*z))/(T_A_K))**5.26)/1.01/(T_A_K)/0.287
        c_p = np.tile(1004.16,np.shape(T_A_K))
    else:
        rho = calc_rho(p, ea, T_A_K)  # Air density
        c_p = calc_c_p(p, ea)  # Heat capacity of air
    z_0H = calc_z_0H(z_0M, kB=kB)  # Roughness length for heat transport

    # Calculate LAI dependent parameters for dataset where LAI > 0
    omega0 = calc_omega0_Kustas(LAI, f_c, x_LAD=x_LAD, isLAIeff=True)
    F = np.asarray(LAI / f_c)  # Real LAI
    # Fraction of vegetation observed by the sensor
    f_theta = calc_F_theta_campbell(vza, F, w_C=w_C, Omega0=omega0, x_LAD=x_LAD)

    # Initially assume stable atmospheric conditions and set variables for
    # iteration of the Monin-Obukhov length
    u_friction = calc_u_star(u, z_u, L, d_0, z_0M)
    u_friction = np.asarray(np.maximum(u_friction_min, u_friction))
    L_old = np.ones(Tr_K.shape)
    L_diff = np.asarray(np.ones(Tr_K.shape))
    # First assume that canopy temperature equals the minumum of Air or
    # radiometric T
    T_C = np.asarray(np.minimum(Tr_K, T_A_K))
    flag, T_S = calc_T_S(Tr_K, T_C, f_theta)
    flag[nullMask==-9999]=255
    iterMask = np.zeros(LAI.shape)
    conv1 = 0.0
    # Outer loop for estimating stability.
    # Stops when difference in consecutives L is below a given threshold
    for n_iterations in range(max_iterations):
        i = flag != 255
        #if np.all(L_diff[i] < L_thres):
        if conv1 > 0.98:
            #print("Finished interation with a max. L diff: " + str(np.max(L_diff)))
            break
#        print("Iteration " + str(n_iterations) +
#              ", max. L diff: " + str(np.max(L_diff)))
        iterations[
            np.logical_and(
                L_diff >= L_thres,
                flag != 255)] = n_iterations

        # Inner loop to iterativelly reduce alpha_PT in case latent heat flux
        # from the soil is negative. The initial assumption is of potential
        # canopy transpiration.
        flag[np.logical_and(L_diff >= L_thres, flag != 255)] = 0
        LE_S[np.logical_and(L_diff >= L_thres, flag != 255)] = -1
        alpha_PT_rec = np.asarray(alpha_PT + 0.1)
        while np.any(LE_S[i] < 0):
            i = np.logical_and.reduce(
                (LE_S < 0, L_diff >= L_thres, flag != 255))

            alpha_PT_rec[i] -= 0.1

            # There cannot be negative transpiration from the vegetation
            alpha_PT_rec[alpha_PT_rec <= 0.0] = 0.0
            flag[np.logical_and(i, alpha_PT_rec == 0.0)] = 5

            flag[
                np.logical_and.reduce(
                    (i, alpha_PT_rec < alpha_PT, alpha_PT_rec > 0.0))] = 3

            # Calculate the aerodynamic resistance
            R_A[i] = calc_R_A(z_T[i], u_friction[i], L[i], d_0[i], z_0H[i])
            # Calculate soil and canopy resistances
            U_C = calc_u_C_star(
                u_friction[i], hc[i], d_0[i], z_0M[i], L=L[i])
            if resistance_form == 0:
                # Wind speed is highly attenuated within the canopy volume
                u_d_zm = calc_u_Goudriaan(
                    U_C, hc[i], F[i], leaf_width[i], d_0[i] + z_0M[i])

                # Vegetation in series with soil, i.e. well mixed, so we use
                # the landscape LAI
                R_x[i] = calc_R_x_Norman(LAI[i], leaf_width[i], u_d_zm)
                i = np.logical_and.reduce(
                (LE_S < 0, flag != 255, L_diff >= L_thres,
                 np.logical_not(np.isnan(R_x))))
                # Calculate soil and canopy resistances
                U_C = calc_u_C_star(
                u_friction[i], hc[i], d_0[i], z_0M[i], L=L[i])
                # Clumped vegetation enhanced wind speed for the soil surface
                u_S = calc_u_Goudriaan(
                    U_C, hc[i], omega0[i] * F[i], leaf_width[i], z0_soil[i])
                R_S[i] = calc_R_S_Kustas(u_S, T_S[i] - T_C[i])
            elif resistance_form == 1:
                # Vegetation in series with soil, i.e. well mixed, so we use
                # the landscape LAI
                R_x[i] = calc_R_x_Choudhury(U_C, LAI[i], leaf_width[i])
                R_S[i] = calc_R_S_Choudhury(
                    u_friction[i], hc[i], z_0M[i], d_0[i], z_u[i], z0_soil[i])
            elif resistance_form == 2:
                # Vegetation in series with soil, i.e. well mixed, so we use
                # the landscape LAI
                R_x[i] = calc_R_x_McNaughton(
                    LAI[i], leaf_width[i], u_friction[i])
                R_S[i] = calc_R_S_McNaughton(u_friction[i])
            elif resistance_form == 3:
                # Clumped vegetation enhanced wind speed for the soil surface
                alpha_k = calc_A_Goudriaan(
                    hc[i], omega0[i] * F[i], leaf_width[i])
                # Wind speed is highly attenuated within the canopy volume
                alpha_prime = calc_A_Goudriaan(hc[i], F[i], leaf_width[i])
                # Vegetation in series with soil, i.e. well mixed, so we use
                # the landscape LAI
                R_x[i] = calc_R_x_Choudhury(
                    U_C, LAI[i], leaf_width[i], alpha_prime=alpha_prime)
                R_S[i] = calc_R_S_Choudhury(u_friction[i], hc[i], z_0M[i], d_0[
                                               i], z_u, z0_soil[i], alpha_k=alpha_k)
            else:
                # Clumped vegetation enhanced wind speed for the soil surface
                u_S = calc_u_Goudriaan(
                    U_C, hc[i], omega0[i] * F[i], leaf_width[i], z0_soil[i])
                # Wind speed is highly attenuated within the canopy volume
                u_d_zm = calc_u_Goudriaan(
                    U_C, hc[i], F[i], leaf_width[i], d_0[i] + z_0M[i])
                # Vegetation in series with soil, i.e. well mixed, so we use
                # the landscape LAI
                R_x[i] = calc_R_x_Norman(LAI[i], leaf_width[i], u_d_zm)
                R_S[i] = calc_R_S_Kustas(u_S, T_S[i] - T_C[i])
            R_S = np.asarray(np.maximum(1e-3, R_S))
            R_x = np.asarray(np.maximum(1e-3, R_x))
            R_A = np.asarray(np.maximum(1e-3, R_A))

            # Calculate net longwave radiation with current values of T_C and T_S
            Ln_C[i], Ln_S[i] = calc_L_n_Kustas(
                T_C[i], T_S[i], L_dn[i], LAI[i], emis_C[i], emis_S[i])
            delta_Rn = Sn_C + Ln_C
            Rn_S = Sn_S + Ln_S
            # Calculate the canopy and soil temperatures using the Priestley
            # Taylor appoach
            H_C[i] = calc_H_C_PT(
                delta_Rn[i],
                f_g[i],
                T_A_K[i],
                p[i],
                c_p[i],
                alpha_PT_rec[i])
            T_C[i] = calc_T_C_series(Tr_K[i], T_A_K[i], R_A[i], R_x[i], R_S[
                                   i], f_theta[i], H_C[i], rho[i], c_p[i])

            # Calculate soil temperature
            flag_t = np.zeros(flag.shape)
            flag_t[i], T_S[i] = calc_T_S(Tr_K[i], T_C[i], f_theta[i])
            flag[flag_t == 255] = 255
            LE_S[flag_t == 255] = 0

            # Recalculate soil resistance using new soil temperature
            if resistance_form == 0:
                R_S[i] = calc_R_S_Kustas(u_S, T_S[i] - T_C[i])
                R_S = np.asarray(np.maximum(1e-3, R_S))

            i = np.logical_and.reduce(
                (LE_S < 0, flag != 255, L_diff >= L_thres,
                 np.logical_not(np.isnan(R_x))))

            # Get air temperature at canopy interface
            T_AC[i] = ((T_A_K[i] / R_A[i] + T_S[i] / R_S[i] + T_C[i] / R_x[i])
                       / (1.0 / R_A[i] + 1.0 / R_S[i] + 1.0 / R_x[i]))

            # Calculate soil fluxes
            H_S[i] = rho[i] * c_p[i] * (T_S[i] - T_AC[i]) / R_S[i]

            # Compute Soil Heat Flux Ratio
            #G[i] = calc_G([calcG_array[0], calcG_array], Rn_S, i)
            G[i]=calc_G_ratio(Rn_S[i], calcG_array[1][i])

            # Estimate latent heat fluxes as residual of energy balance at the
            # soil and the canopy
            LE_S[i] = Rn_S[i] - G[i] - H_S[i]
            LE_C[i] = delta_Rn[i] - H_C[i]

            # Special case if there is no transpiration from vegetation.
            # In that case, there should also be no evaporation from the soil
            # and the energy at the soil should be conserved.
            # See end of appendix A1 in Guzinski et al. (2015).
            noT = np.logical_and(i, LE_C == 0)
            H_S[noT] = np.minimum(H_S[noT], Rn_S[noT] - G[noT])
            G[noT] = np.maximum(G[noT], Rn_S[noT] - H_S[noT])
            LE_S[noT] = 0

            # Calculate total fluxes
            H = np.asarray(H_C + H_S)
            LE = np.asarray(LE_C + LE_S)
            # Now L can be recalculated and the difference between iterations
            # derived
            if isinstance(UseL, bool):
                L[i] = calc_L(
                    u_friction[i],
                    T_A_K[i],
                    rho[i],
                    c_p[i],
                    H[i],
                    LE[i])
                # Calculate again the friction velocity with the new stability
                # correctios
                u_friction[i] = calc_u_star(
                    u[i], z_u[i], L[i], d_0[i], z_0M[i])
                u_friction = np.asarray(np.maximum(u_friction_min, u_friction))

        if isinstance(UseL, bool):
            L_diff = np.asarray(np.fabs(L - L_old) / np.fabs(L_old))
            #L_diff[np.isnan(L_diff)] = float('inf')
            L_old = np.array(L)
            L_old[L_old == 0] = 1e-36
            iterMask[np.where(L_diff < L_thres)]=1.0
            conv1 = np.sum(iterMask)/(LAI.size)
            #print 'TSEB convergence=%f percent' % (conv1*100.)

    (flag,
     T_S,
     T_C,
     T_AC,
     L_nS,
     L_nC,
     LE_C,
     H_C,
     LE_S,
     H_S,
     G,
     R_S,
     R_x,
     R_A,
     u_friction,
     L,
     n_iterations) = map(np.asarray,
                         (flag,
                          T_S,
                          T_C,
                          T_AC,
                          Ln_S,
                          Ln_C,
                          LE_C,
                          H_C,
                          LE_S,
                          H_S,
                          G,
                          R_S,
                          R_x,
                          R_A,
                          u_friction,
                          L,
                          iterations))

    return flag, T_S, T_C, T_AC, L_nS, L_nC, LE_C, H_C, LE_S, H_S, G, R_S, R_x, R_A, u_friction, L, n_iterations

def calc_F_theta_campbell(theta, F, w_C=1, Omega0=1, x_LAD=1):
    '''Calculates the fraction of vegetatinon observed at an angle.
    Parameters
    ----------
    theta : float
        Angle of incidence (degrees).
    F : float
        Real Leaf (Plant) Area Index.
    w_C : float
        Ratio of vegetation height versus width, optional (default = 1).
    Omega0 : float
        Clumping index at nadir, optional (default =1).
    x_LAD : float
        Chi parameter for the ellipsoidal Leaf Angle Distribution function,
        use x_LAD=1 for a spherical LAD.
    Returns
    -------
    f_theta : float
        fraction of vegetation obsserved at an angle.
    References
    ----------
    .. [Campbell1998] Campbell, G. S. & Norman, J. M. (1998), An introduction to environmental
        biophysics. Springer, New York
        https://archive.org/details/AnIntroductionToEnvironmentalBiophysics.
    .. [Norman1995] J.M. Norman, W.P. Kustas, K.S. Humes, Source approach for estimating
        soil and vegetation energy fluxes in observations of directional radiometric
        surface temperature, Agricultural and Forest Meteorology, Volume 77, Issues 3-4,
        Pages 263-293, http://dx.doi.org/10.1016/0168-1923(95)02265-Y.
    '''

    # First calcualte the angular clumping factor Omega based on eq (3) from
    # W.P. Kustas, J.M. Norman,  Agricultural and Forest Meteorology 94 (1999)
    # CHECK: should theta here be in degrees or radians
    OmegaTheta = Omega0 / (Omega0 + (1.0 - Omega0) *
                           np.exp(-2.2 * np.radians(theta)**(3.8 - 0.46 * w_C)))
    # Estimate the beam extinction coefficient based on a elipsoidal LAD function
    # Eq. 15.4 of Campbell and Norman (1998)
    K_be = calc_K_be_Campbell(theta, x_LAD)
    ftheta = 1.0 - np.exp(-K_be * OmegaTheta * F)
    return np.asarray(ftheta)

def calc_G(calcG_params, Rn_S, i=None):

    if i is None:
        i = np.ones(Rn_S.shape, dtype=bool)
    if calcG_params[0][0] == 0:
        G = calcG_params[1][i]
    elif calcG_params[0][0] == 1:
        G = calc_G_ratio(Rn_S[i], calcG_params[1][i])
    elif calcG_params[0][0] == 2:
        G = calc_G_time_diff(Rn_S[i], [calcG_params[1][i], calcG_params[
                           0][1], calcG_params[0][2], calcG_params[0][3]])
    return np.asarray(G)

def calc_G_time_diff(R_n, G_param=[12.0, 0.35, 3.0, 24.0]):
    ''' Estimates Soil Heat Flux as function of time and net radiation.
    Parameters
    ----------
    R_n : float
        Net radiation (W m-2).
    G_param : tuple(float,float,float,float)
        tuple with parameters required (time, Amplitude,phase_shift,shape).
            time: float
                time of interest (decimal hours).
            Amplitude : float
                maximum value of G/Rn, amplitude, default=0.35.
            phase_shift : float
                shift of peak G relative to solar noon (default 3hrs after noon).
            shape : float
                shape of G/Rn, default 24 hrs.
    Returns
    -------
    G : float
        Soil heat flux (W m-2).
    References
    ----------
    .. [Santanello2003] Joseph A. Santanello Jr. and Mark A. Friedl, 2003: Diurnal Covariation in
        Soil Heat Flux and Net Radiation. J. Appl. Meteor., 42, 851-862,
        http://dx.doi.org/10.1175/1520-0450(2003)042<0851:DCISHF>2.0.CO;2.'''

    # Get parameters
    time = 12.0 - G_param[0]
    A = G_param[1]
    phase_shift = G_param[2]
    B = G_param[3]
    G_ratio = A * np.cos(2.0 * np.pi * (time + phase_shift) / B)
    G = R_n * G_ratio
    return np.asarray(G)

def calc_G_ratio(Rn_S, G_ratio=0.35):
    '''Estimates Soil Heat Flux as ratio of net soil radiation.
    Parameters
    ----------
    Rn_S : float
        Net soil radiation (W m-2).
    G_ratio : float, optional
        G/Rn_S ratio, default=0.35.
    Returns
    -------
    G : float
        Soil heat flux (W m-2).
    References
    ----------
    .. [Choudhury1987] B.J. Choudhury, S.B. Idso, R.J. Reginato, Analysis of an empirical model
        for soil heat flux under a growing wheat crop for estimating evaporation by an
        infrared-temperature based energy balance equation, Agricultural and Forest Meteorology,
        Volume 39, Issue 4, 1987, Pages 283-297,
        http://dx.doi.org/10.1016/0168-1923(87)90021-9.
    '''

    G = G_ratio * Rn_S
    return np.asarray(G)

def calc_H_C_PT(delta_R_ni, f_g, T_A_K, P, c_p, alpha):
    '''Calculates canopy sensible heat flux based on the Priestley and Taylor formula.
    Parameters
    ----------
    delta_R_ni : float
        net radiation divergence of the vegetative canopy (W m-2).
    f_g : float
        fraction of vegetative canopy that is green.
    T_A_K : float
        air temperature (Kelvin).
    P : float
        air pressure (mb).
    c_p : float
        heat capacity of moist air (J kg-1 K-1).
    alpha : float
        the Priestley Taylor parameter.
    Returns
    -------
    H_C : float
        Canopy sensible heat flux (W m-2).
    References
    ----------
    Equation 14 in [Norman1995]_
    '''

    # slope of the saturation pressure curve (kPa./deg C)
    s = calc_delta_vapor_pressure(T_A_K)
    s = s * 10  # to mb
    # latent heat of vaporisation (MJ./kg)
    Lambda = calc_lambda(T_A_K)
    # psychrometric constant (mb C-1)
    gama = calc_psicr(P, Lambda)
    s_gama = s / (s + gama)
    H_C = delta_R_ni * (1.0 - alpha * f_g * s_gama)
    return np.asarray(H_C)

def calc_T_C_series(Tr_K, T_A_K, R_A, R_x, R_S, f_theta, H_C, rho, c_p):
    '''Estimates canopy temperature from canopy sensible heat flux and
    resistance network in series.
    Parameters
    ----------
    Tr_K : float
        Directional Radiometric Temperature (K).
    T_A_K : float
        Air Temperature (K).
    R_A : float
        Aerodynamic resistance to heat transport (s m-1).
    R_x : float
        Bulk aerodynamic resistance to heat transport at the canopy boundary layer (s m-1).
    R_S : float
        Aerodynamic resistance to heat transport at the soil boundary layer (s m-1).
    f_theta : float
        Fraction of vegetation observed.
    H_C : float
        Sensible heat flux of the canopy (W m-2).
    rho : float
        Density of air (km m-3).
    c_p : float
        Heat capacity of air at constant pressure (J kg-1 K-1).
    Returns
    -------
    T_C : float
        Canopy temperature (K).
    References
    ----------
    Eqs. A5-A13 in [Norman1995]_'''

    T_R_K_4 = Tr_K**4
    # equation A7 from Norman 1995, linear approximation of temperature of the
    # canopy
    T_C_lin = ((T_A_K / R_A + Tr_K / (R_S * (1.0 - f_theta))
                + H_C * R_x / (rho * c_p) * (1.0 / R_A + 1.0 / R_S + 1.0 / R_x))
               / (1.0 / R_A + 1.0 / R_S + f_theta / (R_S * (1.0 - f_theta))))
    # equation A12 from Norman 1995
    T_D = (T_C_lin * (1 + R_S / R_A) - H_C * R_x / (rho * c_p)
           * (1.0 + R_S / R_x + R_S / R_A) - T_A_K * R_S / R_A)
    # equation A11 from Norman 1995
    delta_T_C = ((T_R_K_4 - f_theta * T_C_lin**4 - (1.0 - f_theta) * T_D**4) / \
                 (4.0 * (1.0 - f_theta) * T_D**3 * (1.0 + R_S / R_A) + 4.0 * f_theta * T_C_lin**3))
    # get canopy temperature in Kelvin
    T_C = T_C_lin + delta_T_C
    return np.asarray(T_C)

def calc_T_S(T_R, T_C, f_theta):
    '''Estimates soil temperature from the directional LST.
    Parameters
    ----------
    T_R : float
        Directional Radiometric Temperature (K).
    T_C : float
        Canopy Temperature (K).
    f_theta : float
        Fraction of vegetation observed.
    Returns
    -------
    flag : float
        Error flag if inversion not possible (255).
    T_S: float
        Soil temperature (K).
    References
    ----------
    Eq. 1 in [Norman1995]_'''

    # Convert the input scalars to numpy arrays
    T_R, T_C, f_theta = map(np.asarray, (T_R, T_C, f_theta))
    T_temp = T_R**4 - f_theta * T_C**4
    T_S = np.zeros(T_R.shape)
    flag = np.zeros(T_R.shape)

    # Succesfull inversion
    T_S[T_temp >= 0] = (T_temp[T_temp >= 0] /
                        (1.0 - f_theta[T_temp >= 0]))**0.25

    # Unsuccesfull inversion
    T_S[T_temp < 0] = 1e-6
    flag[T_temp < 0] = 255

    return np.asarray(flag), np.asarray(T_S)

def _check_default_parameter_size(parameter, input_array):

    parameter = np.asarray(parameter)
    if parameter.size == 1:
        parameter = np.ones(input_array.shape) * parameter
        return np.asarray(parameter)
    elif parameter.shape != input_array.shape:
        raise ValueError(
            'dimension mismatch between parameter array and input array with shapes %s and %s' %
            (parameter.shape, input_array.shape))
    else:
        return np.asarray(parameter)