#!/usr/bin/env python

"""@package docstring
File: me_solver.py
Author: Adam Lamson
Email: adam.lamson@colorado.edu
Description:
"""

import time
import numpy as np
from scipy.integrate import solve_ivp
from .choose_me_evolver import choose_me_evolver
from .solver import Solver
from .non_dimensionalizer import NonDimensionalizer
from .rod_steric_forces import get_min_dist_vec

class MomentExpansionSolver(Solver):

    """!Solve the evolution of two rods by expanding the Fokker - Planck equation
        in a series of moments of motor end positions on rods.
    """

    def __init__(self, pfile=None, pdict=None):
        """!Set parameters for ODE to be solved including initial conditions.

        @param pfile: yaml parameter file name
        @param pdict: parameter dictionary
        """
        print("Init MomentExpansionSolver ->", end=" ")
        Solver.__init__(self, pfile, pdict)

    def ParseParams(self):
        """!Collect parameters from yaml file or dictionary then calculate
        some necessary parameters if not defined. Also non-dimensionalize parameters.
        @return: void
        """
        Solver.ParseParams(self)

        self.dt = self._params["dt"]  # Time step
        if "nt" not in self._params:
            print("!!! Warning: nt not defined. Using nsteps and dt to find ",
                  "total time. nsteps and dt are not used in this calculation.")
            self.nsteps = int(self._params["nsteps"])
            self.nt = self.nsteps * self.dt
            self._params["nt"] = self.nt
        else:
            self.nt = self._params["nt"]

        if "nwrite" not in self._params:
            self.twrite = self._params["twrite"]
            self.nwrite = int(self.twrite / self.dt)
        elif "twrite" not in self._params:
            self.nwrite = self._params["nwrite"]
            self.twrite = float(self.nwrite * self.dt)
        else:
            print("!!! Warning: Write parameters over defined,",
                  "using twrite to calculate number of steps between write out.")
            self.twrite = self._params["twrite"]
            self.nwrite = int(self.twrite / self.dt)

        self.non_dimmer = self.non_dimensionalize()
        # Rod orientation vectors
        self.R1_vec = np.asarray(self._params['R1_vec'])
        self.R2_vec = np.asarray(self._params['R2_vec'])
        # Make sure to renormalize
        self.R1_vec /= np.linalg.norm(self.R1_vec)
        self.R2_vec /= np.linalg.norm(self.R2_vec)

        # Check to see if steric forces should be used
        self.steric_flag = self._params.get('steric_interactions', None)
        if self.steric_flag == 'constrained':
            self.constr_vec = get_min_dist_vec(self.R2_pos - self.R1_pos,
                                               self.R1_vec, self.R2_vec)
            self.constr_vec /= np.linalg.norm(self.constr_vec)

        print("R1_vec = ", self.R1_vec)
        print("R2_vec = ", self.R2_vec)

        self.t_eval = np.linspace(0, self.nt, int(self.nt / self.twrite) + 1)
        self._nframes = self.t_eval.size

        # Set integration method for solver
        self.method = self._params.get('method', 'LSODA')
        self._params['method'] = self.method
        print("Solving method = ", self.method)

        # Specify the ODE type
        self.ODE_type = self._params.get('ODE_type', 'zrl')
        self._params['ODE_type'] = self.ODE_type
        print("ODE type = ", self.ODE_type)

    def setInitialConditions(self):
        """!Set the initial conditions for the system of ODEs
        @return: void
        """
        self.sol_init = np.zeros(26)
        # Set all geometric variables
        self.sol_init[:12] = np.concatenate(
            (self.R1_pos, self.R2_pos, self.R1_vec, self.R2_vec))
        print("=== Initial conditions ===")
        print(self.sol_init)
        # TODO Allow for different initial conditions of moments besides zero

        # Set solver once you set initial conditions
        # Add kwargs
        self.ode_solver = choose_me_evolver(self.sol_init, self)

    def makeDataframe(self):
        """!Create data frame to be written out
        @return: TODO
        """
        t_arr = self.non_dimmer.dim_val(self.t_eval, ['time'])
        print("Evaluated at:", t_arr)
        if not self.data_frame_made:
            self._time_dset = self._h5_data.create_dataset('time',
                                                           data=t_arr,
                                                           dtype=np.float32)
            self._xl_grp = self._h5_data.create_group('xl_data')
            self._rod_grp = self._h5_data.create_group('rod_data')

            Solver.makeDataframe(self)
            self.data_frame_made = True

    def Run(self):
        """!Run algorithm to solve system of ODEs
        @return: TODO
        """

        t0 = time.time()
        self.sol = solve_ivp(self.ode_solver, [0, self.nt], self.sol_init,
                             t_eval=self.t_eval, method=self.method,)
        # min_step=self.dt, atol=1e-6)
        self.cpu_time = time.time() - t0
        print(
            r" --- Total simulation time {:.4f} seconds ---".format(self.cpu_time))

        self.Write()

    def make_rod_dataset(self):
        """!Initialize dataframe with empty rod configuration data
        @return: void

        """
        self._R1_pos_dset = self._rod_grp.create_dataset(
            'R1_pos', data=self.sol.y[: 3, :].T)
        self._R2_pos_dset = self._rod_grp.create_dataset(
            'R2_pos', data=self.sol.y[3: 6, :].T)
        self._R1_vec_dset = self._rod_grp.create_dataset(
            'R1_vec', data=self.sol.y[6: 9, :].T)
        self._R2_vec_dset = self._rod_grp.create_dataset(
            'R2_vec', data=self.sol.y[9: 12, :].T)

    def make_xl_moment_dataset(self):
        """!Initialize dataframe with empty crosslinker moment data
        @return: void

        """
        self._mu0_dset = self._xl_grp.create_dataset('zeroth_moment',
                                                     data=self.sol.y[12, :].T,
                                                     dtype=np.float32)
        self._mu1_dset = self._xl_grp.create_dataset('first_moments',
                                                     data=self.sol.y[13: 15, :].T,
                                                     dtype=np.float32)
        self._mu2_dset = self._xl_grp.create_dataset('second_moments',
                                                     data=self.sol.y[15:18, :].T,
                                                     dtype=np.float32)
        self._B0_dset = self._xl_grp.create_dataset('zeroth_boundary_terms',
                                                    data=self.sol.y[18:20, :].T,
                                                    dtype=np.float32)
        self._B1_dset = self._xl_grp.create_dataset('first_boundary_terms',
                                                    data=self.sol.y[20:22, :].T,
                                                    dtype=np.float32)
        self._B2_dset = self._xl_grp.create_dataset('second_boundary_terms',
                                                    data=self.sol.y[22:24, :].T,
                                                    dtype=np.float32)
        self._B3_dset = self._xl_grp.create_dataset('third_boundary_terms',
                                                    data=self.sol.y[24:26, :].T,
                                                    dtype=np.float32)

    def Write(self):
        """!Write out data
        @return: void

        """
        self.redimensionalize()
        self.make_xl_moment_dataset()
        self.make_rod_dataset()
        # Store how long the simulation took
        self._h5_data.attrs['cpu_time'] = self.cpu_time

    def non_dimensionalize(self):
        """!Non-dimensionalize parameters to reduce error in calculations.
        @return: non dimensionalizer

        """
        # non_dim_dict = {'time': 1. / self._params['ko'],
        #                 # 'length': max(self._params['L1'], self._params['L2']),
        #                 'length': self._params['fs'] / self._params['ks'],
        #                 'energy': 1. / self._params['beta']}
        # NonDimensionalizer not working currently
        non_dim_dict = {'time': 1.,
                        'length': float(max(self._params['L1'],
                                            self._params['L2'])),
                        # 'length': 1.,
                        'energy': 1.}
        non_dimmer = NonDimensionalizer(**non_dim_dict)
        # non_dimmer.calc_new_dim('force', ['energy', 'length'], [1, -1])

        self.beta = non_dimmer.non_dim_val(self._params['beta'],
                                           ['energy'], [-1])
        self.visc = non_dimmer.non_dim_val(self._params['viscosity'],
                                           ['energy', 'time', 'length'],
                                           [1, 1, -3])
        self.L_i = non_dimmer.non_dim_val(self._params['L1'], ['length'])
        self.L_j = non_dimmer.non_dim_val(self._params['L2'], ['length'])
        self.R1_pos = non_dimmer.non_dim_val(
            self._params['R1_pos'], ['length'])
        self.R2_pos = non_dimmer.non_dim_val(
            self._params['R2_pos'], ['length'])
        self.rod_diam = non_dimmer.non_dim_val(self._params['rod_diameter'],
                                               ['length'])
        self.dt = non_dimmer.non_dim_val(self.dt, ['time'])
        self.nt = non_dimmer.non_dim_val(self.nt, ['time'])
        self.twrite = non_dimmer.non_dim_val(self.twrite, ['time'])
        self.ko = non_dimmer.non_dim_val(self._params['ko'], ['time'], [-1])
        self.co = non_dimmer.non_dim_val(self._params['co'], ['length'], [-2])
        self.ks = non_dimmer.non_dim_val(self._params['ks'],
                                         ['energy', 'length'], [1, -2])
        self.ho = non_dimmer.non_dim_val(self._params['ho'], ['length'])
        self.vo = non_dimmer.non_dim_val(self._params['vo'],
                                         ['length', 'time'], [1, -1])
        self.fs = non_dimmer.non_dim_val(
            self._params['fs'], ['energy', 'length'], [1, -1])
        return non_dimmer

    def redimensionalize(self):
        """!Redimensionalize data arrays
        @return: void

        """
        # Redimensionalize rod positions
        self.sol.y[:6, :] = self.non_dimmer.dim_val(self.sol.y[:6, :],
                                                    ['length'])
        # Redimensionalize first moments
        self.sol.y[13:15, :] = self.non_dimmer.dim_val(self.sol.y[13: 15, :],
                                                       ['length'])
        # Redimensionalize second moments
        self.sol.y[15:18, :] = self.non_dimmer.dim_val(self.sol.y[15:18, :],
                                                       ['length'], [2])
        self.sol.y[20:22, :] = self.non_dimmer.dim_val(self.sol.y[20:22, :],
                                                       ['length'])
        self.sol.y[22:24, :] = self.non_dimmer.dim_val(self.sol.y[22:24, :],
                                                       ['length'], [2])
        self.sol.y[24:26, :] = self.non_dimmer.dim_val(self.sol.y[24:26, :],
                                                       ['length'], [3])


import numpy as np

try:
    import classy
    HAVE_CLASS = True
except ImportError:
    HAVE_CLASS = False

try:
    import camb
    import camb.model
    HAVE_CAMB = True
except ImportError:
    HAVE_CAMB = False

try:
    import isitgr  # noqa: F401
except ImportError:
    pass  # prevent nans from isitgr

from . import ccllib as lib
from .pyutils import check
from .pk2d import Pk2D
from .errors import CCLError

def get_camb_pk_lin(cosmo, nonlin=False):
    """Run CAMB and return the linear power spectrum.

    Args:
        cosmo (:class:`~pyccl.core.Cosmology`): Cosmological
            parameters. The cosmological parameters with
            which to run CAMB.
        nonlin (:obj:`bool`, optional): Whether to compute and return the
            non-linear power spectrum as well.

    Returns:
        :class:`~pyccl.pk2d.Pk2D`: Power spectrum object. The linear power \
            spectrum. If ``nonlin=True``, returns a tuple \
            ``(pk_lin, pk_nonlin)``.
    """

    # Comment from Jarvis: TODO clean up this and other assert
    # anti-patterns in this file
    assert HAVE_CAMB, (
        "You must have the `camb` python package "
        "installed to run CCL with CAMB!")

    # Get extra CAMB parameters that were specified
    extra_camb_params = {}
    try:
        extra_camb_params = cosmo["extra_parameters"]["camb"]
    except (KeyError, TypeError):
        pass

    # z sampling from CCL parameters
    na = lib.get_pk_spline_na(cosmo.cosmo)
    status = 0
    a_arr, status = lib.get_pk_spline_a(cosmo.cosmo, na, status)
    check(status)
    a_arr = np.sort(a_arr)
    zs = 1.0 / a_arr - 1
    zs = np.clip(zs, 0, np.inf)

    # deal with normalization
    if np.isfinite(cosmo["A_s"]):
        A_s_fid = cosmo["A_s"]
    elif np.isfinite(cosmo["sigma8"]):
        # in this case, CCL will internally normalize for us when we init
        # the linear power spectrum - so we just get close
        A_s_fid = 2.43e-9 * (cosmo["sigma8"] / 0.87659)**2
    else:
        raise CCLError(
            "Could not normalize the linear power spectrum! "
            "A_s = %f, sigma8 = %f" % (
                cosmo['A_s'], cosmo['sigma8']))

    # init camb params
    cp = camb.model.CAMBparams()

    # turn some stuff off
    cp.WantCls = False
    cp.DoLensing = False
    cp.Want_CMB = False
    cp.Want_CMB_lensing = False
    cp.Want_cl_2D_array = False
    cp.WantTransfer = True

    # basic background stuff
    h2 = cosmo['h']**2
    cp.H0 = cosmo['h'] * 100
    cp.ombh2 = cosmo['Omega_b'] * h2
    cp.omch2 = cosmo['Omega_c'] * h2
    cp.omk = cosmo['Omega_k']

    # "constants"
    cp.TCMB = cosmo['T_CMB']

    # neutrinos
    # We maually setup the CAMB neutrinos to match the adjustments CLASS
    # makes to their temperatures.
    cp.share_delta_neff = False
    cp.omnuh2 = cosmo['Omega_nu_mass'] * h2
    cp.num_nu_massless = cosmo['N_nu_rel']
    cp.num_nu_massive = int(cosmo['N_nu_mass'])
    cp.nu_mass_eigenstates = int(cosmo['N_nu_mass'])

    delta_neff = cosmo['Neff'] - 3.046  # used for BBN YHe comps

    # CAMB defines a neutrino degeneracy factor as T_i = g^(1/4)*T_nu
    # where T_nu is the standard neutrino temperature from first order
    # computations
    # CLASS defines the temperature of each neutrino species to be
    # T_i_eff = TNCDM * T_cmb where TNCDM is a fudge factor to get the
    # total mass in terms of eV to match second-order computations of the
    # relationship between m_nu and Omega_nu.
    # We are trying to get both codes to use the same neutrino temperature.
    # thus we set T_i_eff = T_i = g^(1/4) * T_nu and solve for the right
    # value of g for CAMB. We get g = (TNCDM / (11/4)^(-1/3))^4
    g = np.power(
        lib.cvar.constants.TNCDM / np.power(11.0/4.0, -1.0/3.0),
        4.0)

    if cosmo['N_nu_mass'] > 0:
        nu_mass_fracs = cosmo['m_nu'][:cosmo['N_nu_mass']]
        nu_mass_fracs = nu_mass_fracs / np.sum(nu_mass_fracs)

        cp.nu_mass_numbers = np.ones(cosmo['N_nu_mass'], dtype=np.int)
        cp.nu_mass_fractions = nu_mass_fracs
        cp.nu_mass_degeneracies = np.ones(int(cosmo['N_nu_mass'])) * g
    else:
        cp.nu_mass_numbers = []
        cp.nu_mass_fractions = []
        cp.nu_mass_degeneracies = []

    # get YHe from BBN
    cp.bbn_predictor = camb.bbn.get_predictor()
    cp.YHe = cp.bbn_predictor.Y_He(
        cp.ombh2 * (camb.constants.COBE_CMBTemp / cp.TCMB) ** 3,
        delta_neff)

    camb_de_models = ['DarkEnergyPPF', 'ppf', 'DarkEnergyFluid', 'fluid']
    camb_de_model = extra_camb_params.get('dark_energy_model', 'fluid')
    if camb_de_model not in camb_de_models:
        raise ValueError("The only dark energy models CCL supports with"
                         " camb are fluid and ppf.")
    cp.set_classes(
        dark_energy_model=camb_de_model
    )

    if camb_de_model not in camb_de_models[:2] and cosmo['wa'] and \
            (cosmo['w0'] < -1 - 1e-6 or
                1 + cosmo['w0'] + cosmo['wa'] < - 1e-6):
        raise ValueError("If you want to use w crossing -1,"
                         " then please set the dark_energy_model to ppf.")
    cp.DarkEnergy.set_params(
        w=cosmo['w0'],
        wa=cosmo['wa']
    )

    if nonlin:
        cp.NonLinearModel = camb.nonlinear.Halofit()
        halofit_version = extra_camb_params.get("halofit_version", "mead")
        options = {k: extra_camb_params[k] for k in
                   ["HMCode_A_baryon",
                    "HMCode_eta_baryon",
                    "HMCode_logT_AGN"] if k in extra_camb_params}
        cp.NonLinearModel.set_params(halofit_version=halofit_version,
                                     **options)

    cp.set_matter_power(
        redshifts=[_z for _z in zs],
        kmax=extra_camb_params.get("kmax", 10.0),
        nonlinear=nonlin)
    if not nonlin:
        assert cp.NonLinear == camb.model.NonLinear_none

    cp.set_for_lmax(extra_camb_params.get("lmax", 5000))
    cp.InitPower.set_params(
        As=A_s_fid,
        ns=cosmo['n_s'])

    # run CAMB and get results
    camb_res = camb.get_results(cp)
    k, z, pk = camb_res.get_linear_matter_power_spectrum(
        hubble_units=True, nonlinear=False)

    # convert to non-h inverse units
    k *= cosmo['h']
    pk /= (h2 * cosmo['h'])

    # now build interpolant
    nk = k.shape[0]
    lk_arr = np.log(k)
    a_arr = 1.0 / (1.0 + z)
    na = a_arr.shape[0]
    sinds = np.argsort(a_arr)
    a_arr = a_arr[sinds]
    ln_p_k_and_z = np.zeros((na, nk), dtype=np.float64)
    for i, sind in enumerate(sinds):
        ln_p_k_and_z[i, :] = np.log(pk[sind, :])

    pk_lin = Pk2D(
        pkfunc=None,
        a_arr=a_arr,
        lk_arr=lk_arr,
        pk_arr=ln_p_k_and_z,
        is_logp=True,
        extrap_order_lok=1,
        extrap_order_hik=2,
        cosmo=cosmo)

    if not nonlin:
        return pk_lin
    else:
        k, z, pk = camb_res.get_linear_matter_power_spectrum(
            hubble_units=True, nonlinear=True)

        # convert to non-h inverse units
        k *= cosmo['h']
        pk /= (h2 * cosmo['h'])

        # now build interpolant
        nk = k.shape[0]
        lk_arr = np.log(k)
        a_arr = 1.0 / (1.0 + z)
        na = a_arr.shape[0]
        sinds = np.argsort(a_arr)
        a_arr = a_arr[sinds]
        ln_p_k_and_z = np.zeros((na, nk), dtype=np.float64)
        for i, sind in enumerate(sinds):
            ln_p_k_and_z[i, :] = np.log(pk[sind, :])

        pk_nonlin = Pk2D(
            pkfunc=None,
            a_arr=a_arr,
            lk_arr=lk_arr,
            pk_arr=ln_p_k_and_z,
            is_logp=True,
            extrap_order_lok=1,
            extrap_order_hik=2,
            cosmo=cosmo)

        return pk_lin, pk_nonlin

def get_isitgr_pk_lin(cosmo):
    """Run ISiTGR-CAMB and return the linear power spectrum.

    Args:
        cosmo (:class:`~pyccl.core.Cosmology`): Cosmological
            parameters. The cosmological parameters with
            which to run ISiTGR-CAMB.

    Returns:
        :class:`~pyccl.pk2d.Pk2D`: Power spectrum \
            object. The linear power spectrum.
    """

    try:
        import isitgr  # noqa: F811
        import isitgr.model
    except ImportError as e:
        e.args = (
            "You must have the `isitgr` python package "
            "installed to run CCL with ISiTGR-CAMB!",
            *e.args)
        raise

    # Get extra CAMB parameters that were specified
    extra_camb_params = {}
    try:
        extra_camb_params = cosmo["extra_parameters"]["camb"]
    except (KeyError, TypeError):
        pass

    # z sampling from CCL parameters
    na = lib.get_pk_spline_na(cosmo.cosmo)
    status = 0
    a_arr, status = lib.get_pk_spline_a(cosmo.cosmo, na, status)
    check(status)
    a_arr = np.sort(a_arr)
    zs = 1.0 / a_arr - 1
    zs = np.clip(zs, 0, np.inf)

    # deal with normalization
    if np.isfinite(cosmo["A_s"]):
        A_s_fid = cosmo["A_s"]
    elif np.isfinite(cosmo["sigma8"]):
        # in this case, CCL will internally normalize for us when we init
        # the linear power spectrum - so we just get close
        A_s_fid = 2.43e-9 * (cosmo["sigma8"] / 0.87659)**2
    else:
        raise CCLError(
            "Could not normalize the linear power spectrum! "
            "A_s = %f, sigma8 = %f" % (
                cosmo['A_s'], cosmo['sigma8']))

    # init isitgr params
    cp = isitgr.model.CAMBparams()

    # turn some stuff off
    cp.WantCls = False
    cp.DoLensing = False
    cp.Want_CMB = False
    cp.Want_CMB_lensing = False
    cp.Want_cl_2D_array = False
    cp.WantTransfer = True

    # basic background stuff
    h2 = cosmo['h']**2
    cp.H0 = cosmo['h'] * 100
    cp.ombh2 = cosmo['Omega_b'] * h2
    cp.omch2 = cosmo['Omega_c'] * h2
    cp.omk = cosmo['Omega_k']
#   cp.GR = 1 means GR modified!
    cp.GR = 1
    cp.ISiTGR_muSigma = True
    cp.mu0 = cosmo['mu_0']
    cp.Sigma0 = cosmo['sigma_0']
    cp.c1 = cosmo['c1_mg']
    cp.c2 = cosmo['c2_mg']
    cp.Lambda = cosmo['lambda_mg']

    # "constants"
    cp.TCMB = cosmo['T_CMB']

    # neutrinos
    # We maually setup the CAMB neutrinos to match the adjustments CLASS
    # makes to their temperatures.
    cp.share_delta_neff = False
    cp.omnuh2 = cosmo['Omega_nu_mass'] * h2
    cp.num_nu_massless = cosmo['N_nu_rel']
    cp.num_nu_massive = int(cosmo['N_nu_mass'])
    cp.nu_mass_eigenstates = int(cosmo['N_nu_mass'])

    delta_neff = cosmo['Neff'] - 3.046  # used for BBN YHe comps

    # ISiTGR built on CAMB which defines a neutrino degeneracy
    # factor as T_i = g^(1/4)*T_nu
    # where T_nu is the standard neutrino temperature from first order
    # computations
    # CLASS defines the temperature of each neutrino species to be
    # T_i_eff = TNCDM * T_cmb where TNCDM is a fudge factor to get the
    # total mass in terms of eV to match second-order computations of the
    # relationship between m_nu and Omega_nu.
    # We are trying to get both codes to use the same neutrino temperature.
    # thus we set T_i_eff = T_i = g^(1/4) * T_nu and solve for the right
    # value of g for CAMB. We get g = (TNCDM / (11/4)^(-1/3))^4
    g = np.power(
        lib.cvar.constants.TNCDM / np.power(11.0/4.0, -1.0/3.0),
        4.0)

    if cosmo['N_nu_mass'] > 0:
        nu_mass_fracs = cosmo['m_nu'][:cosmo['N_nu_mass']]
        nu_mass_fracs = nu_mass_fracs / np.sum(nu_mass_fracs)

        cp.nu_mass_numbers = np.ones(cosmo['N_nu_mass'], dtype=np.int)
        cp.nu_mass_fractions = nu_mass_fracs
        cp.nu_mass_degeneracies = np.ones(int(cosmo['N_nu_mass'])) * g
    else:
        cp.nu_mass_numbers = []
        cp.nu_mass_fractions = []
        cp.nu_mass_degeneracies = []

    # get YHe from BBN
    cp.bbn_predictor = isitgr.bbn.get_predictor()
    cp.YHe = cp.bbn_predictor.Y_He(
        cp.ombh2 * (isitgr.constants.COBE_CMBTemp / cp.TCMB) ** 3,
        delta_neff)

    camb_de_models = ['DarkEnergyPPF', 'ppf', 'DarkEnergyFluid', 'fluid']
    camb_de_model = extra_camb_params.get('dark_energy_model', 'fluid')
    if camb_de_model not in camb_de_models:
        raise ValueError("The only dark energy models CCL supports with"
                         " camb are fluid and ppf.")
    cp.set_classes(
        dark_energy_model=camb_de_model
    )
    if camb_de_model not in camb_de_models[:2] and cosmo['wa'] and \
            (cosmo['w0'] < -1 - 1e-6 or
                1 + cosmo['w0'] + cosmo['wa'] < - 1e-6):
        raise ValueError("If you want to use w crossing -1,"
                         " then please set the dark_energy_model to ppf.")
    cp.DarkEnergy.set_params(
        w=cosmo['w0'],
        wa=cosmo['wa']
    )
    # cp.set_cosmology()
    cp.set_matter_power(
        redshifts=[_z for _z in zs],
        kmax=10,
        nonlinear=False)
    assert cp.NonLinear == isitgr.model.NonLinear_none

    cp.set_for_lmax(5000)
    cp.InitPower.set_params(
        As=A_s_fid,
        ns=cosmo['n_s'])

    # run ISITGR and get results
    isitgr_res = isitgr.get_results(cp)
    k, z, pk = isitgr_res.get_linear_matter_power_spectrum(
        hubble_units=True, nonlinear=False)

    # convert to non-h inverse units
    k *= cosmo['h']
    pk /= (h2 * cosmo['h'])

    # now build interpolant
    nk = k.shape[0]
    lk_arr = np.log(k)
    a_arr = 1.0 / (1.0 + z)
    na = a_arr.shape[0]
    sinds = np.argsort(a_arr)
    a_arr = a_arr[sinds]
    ln_p_k_and_z = np.zeros((na, nk), dtype=np.float64)
    for i, sind in enumerate(sinds):
        ln_p_k_and_z[i, :] = np.log(pk[sind, :])

    pk_lin = Pk2D(
        pkfunc=None,
        a_arr=a_arr,
        lk_arr=lk_arr,
        pk_arr=ln_p_k_and_z,
        is_logp=True,
        extrap_order_lok=1,
        extrap_order_hik=2,
        cosmo=cosmo)
    return pk_lin

def get_class_pk_lin(cosmo):
    """Run CLASS and return the linear power spectrum.

    Args:
        cosmo (:class:`~pyccl.core.Cosmology`): Cosmological
            parameters. The cosmological parameters with
            which to run CLASS.

    Returns:
        :class:`~pyccl.pk2d.Pk2D`: Power spectrum object.\
            The linear power spectrum.
    """

    assert HAVE_CLASS, (
        "You must have the python wrapper for CLASS "
        "installed to run CCL with CLASS!")

    params = {
        "output": "mPk",
        "non linear": "none",
        "P_k_max_1/Mpc": cosmo.cosmo.spline_params.K_MAX_SPLINE,
        "z_max_pk": 1.0/cosmo.cosmo.spline_params.A_SPLINE_MINLOG_PK-1.0,
        "modes": "s",
        "lensing": "no",
        "h": cosmo["h"],
        "Omega_cdm": cosmo["Omega_c"],
        "Omega_b": cosmo["Omega_b"],
        "Omega_k": cosmo["Omega_k"],
        "n_s": cosmo["n_s"]}

    # cosmological constant?
    # set Omega_Lambda = 0.0 if w !=-1 or wa != 0
    if cosmo['w0'] != -1 or cosmo['wa'] != 0:
        params["Omega_Lambda"] = 0
        params['w0_fld'] = cosmo['w0']
        params['wa_fld'] = cosmo['wa']

    # neutrino parameters
    # massless neutrinos
    if cosmo["N_nu_rel"] > 1e-4:
        params["N_ur"] = cosmo["N_nu_rel"]
    else:
        params["N_ur"] = 0.0

    # massive neutrinos
    if cosmo["N_nu_mass"] > 0:
        params["N_ncdm"] = cosmo["N_nu_mass"]
        masses = lib.parameters_get_nu_masses(cosmo._params, 3)
        params["m_ncdm"] = ", ".join(
            ["%g" % m for m in masses[:cosmo["N_nu_mass"]]])

    params["T_cmb"] = cosmo["T_CMB"]

    # if we have sigma8, we need to find A_s
    if np.isfinite(cosmo["A_s"]):
        params["A_s"] = cosmo["A_s"]
    elif np.isfinite(cosmo["sigma8"]):
        # in this case, CCL will internally normalize for us when we init
        # the linear power spectrum - so we just get close
        A_s_fid = 2.43e-9 * (cosmo["sigma8"] / 0.87659)**2
        params["A_s"] = A_s_fid
    else:
        raise CCLError(
            "Could not normalize the linear power spectrum! "
            "A_s = %f, sigma8 = %f" % (
                cosmo['A_s'], cosmo['sigma8']))

    model = None
    try:
        model = classy.Class()
        model.set(params)
        model.compute()

        # Set k and a sampling from CCL parameters
        nk = lib.get_pk_spline_nk(cosmo.cosmo)
        na = lib.get_pk_spline_na(cosmo.cosmo)
        status = 0
        a_arr, status = lib.get_pk_spline_a(cosmo.cosmo, na, status)
        check(status)

        # FIXME - getting the lowest CLASS k value from the python interface
        # appears to be broken - setting to 1e-5 which is close to the
        # old value
        lk_arr = np.log(np.logspace(
            -5,
            np.log10(cosmo.cosmo.spline_params.K_MAX_SPLINE), nk))

        # we need to cut this to the max value used for calling CLASS
        msk = lk_arr < np.log(cosmo.cosmo.spline_params.K_MAX_SPLINE)
        nk = int(np.sum(msk))
        lk_arr = lk_arr[msk]

        # now do interp by hand
        ln_p_k_and_z = np.zeros((na, nk), dtype=np.float64)
        for aind in range(na):
            z = max(1.0 / a_arr[aind] - 1, 1e-10)
            for kind in range(nk):
                ln_p_k_and_z[aind, kind] = np.log(
                    model.pk_lin(np.exp(lk_arr[kind]), z))
    finally:
        if model is not None:
            model.struct_cleanup()
            model.empty()

    params["P_k_max_1/Mpc"] = cosmo.cosmo.spline_params.K_MAX_SPLINE

    # make the Pk2D object
    pk_lin = Pk2D(
        pkfunc=None,
        a_arr=a_arr,
        lk_arr=lk_arr,
        pk_arr=ln_p_k_and_z,
        is_logp=True,
        extrap_order_lok=1,
        extrap_order_hik=2,
        cosmo=cosmo)

    return pk_lin


# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function
import os
import json
import copy
import pprint
import logging
import numpy as np
from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.table import Table, Column
from gammapy.maps import WcsNDMap, MapCoord
import fermipy.config
from fermipy import utils
from fermipy import defaults
from fermipy import wcs_utils
from fermipy import fits_utils
from fermipy.sourcefind_utils import fit_error_ellipse
from fermipy.sourcefind_utils import find_peaks
from fermipy.skymap import Map
from fermipy.config import ConfigSchema
from fermipy.gtutils import FreeParameterState, SourceMapState
from fermipy.timing import Timer
from fermipy.model_utils import get_function_norm_par_name
from LikelihoodState import LikelihoodState
import pyLikelihood as pyLike

class SourceFind(object):
    """Mixin class which provides source-finding functionality to
    `~fermipy.gtanalysis.GTAnalysis`."""

    def find_sources(self, prefix='', **kwargs):
        """An iterative source-finding algorithm that uses likelihood
        ratio (TS) maps of the region of interest to find new sources.
        After each iteration a new TS map is generated incorporating
        sources found in the previous iteration.  The method stops
        when the number of iterations exceeds ``max_iter`` or no
        sources exceeding ``sqrt_ts_threshold`` are found.

        Parameters
        ----------
        {options}

        tsmap : dict
           Keyword arguments dictionary for tsmap method.

        tscube : dict
           Keyword arguments dictionary for tscube method.

        Returns
        -------

        peaks : list
           List of peak objects.

        sources : list
           List of source objects.

        """
        timer = Timer.create(start=True)
        self.logger.info('Starting.')

        schema = ConfigSchema(self.defaults['sourcefind'],
                              tsmap=self.defaults['tsmap'],
                              tscube=self.defaults['tscube'])

        schema.add_option('search_skydir', None, '', SkyCoord)
        schema.add_option('search_minmax_radius', [None, 1.0], '', list)

        config = utils.create_dict(self.config['sourcefind'],
                                   tsmap=self.config['tsmap'],
                                   tscube=self.config['tscube'])
        config = schema.create_config(config, **kwargs)

        # Defining default properties of test source model
        config['model'].setdefault('Index', 2.0)
        config['model'].setdefault('SpectrumType', 'PowerLaw')
        config['model'].setdefault('SpatialModel', 'PointSource')
        config['model'].setdefault('Prefactor', 1E-13)

        o = {'sources': [], 'peaks': []}

        for i in range(config['max_iter']):
            srcs, peaks = self._find_sources_iterate(prefix, i, **config)

            self.logger.info('Found %i sources in iteration %i.' %
                             (len(srcs), i))

            o['sources'] += srcs
            o['peaks'] += peaks
            if len(srcs) == 0:
                break

        self.logger.info('Done.')
        self.logger.info('Execution time: %.2f s', timer.elapsed_time)

        return o

    def _build_src_dicts_from_peaks(self, peaks, maps, src_dict_template):

        tsmap = maps['ts']
        amp = maps['amplitude']

        src_dicts = []
        names = []

        for p in peaks:

            o = fit_error_ellipse(tsmap, (p['ix'], p['iy']), dpix=2)
            skydir = o['skydir']
            p['fit_loc'] = o
            p['fit_skydir'] = o['skydir']

            p.update(o)
            name = utils.create_source_name(skydir)
            src_dict = copy.deepcopy(src_dict_template)
            norm_par = get_function_norm_par_name(
                src_dict_template['SpectrumType'])
            src_dict.update({norm_par: amp.data[p['iy'], p['ix']],
                             'ra': skydir.icrs.ra.deg,
                             'dec': skydir.icrs.dec.deg})

            src_dict['glon_err'] = o['glon_err']
            src_dict['glat_err'] = o['glat_err']
            src_dict['ra_err'] = o['ra_err']
            src_dict['dec_err'] = o['dec_err']
            src_dict['pos_err'] = o['pos_err']
            src_dict['pos_err_semimajor'] = o['pos_err_semimajor']
            src_dict['pos_err_semiminor'] = o['pos_err_semiminor']
            src_dict['pos_r68'] = o['pos_r68']
            src_dict['pos_r95'] = o['pos_r95']
            src_dict['pos_r99'] = o['pos_r99']
            src_dict['pos_angle'] = o['pos_angle']
            src_dict['pos_gal_cov'] = o['pos_gal_cov']
            src_dict['pos_gal_corr'] = o['pos_gal_corr']
            src_dict['pos_cel_cov'] = o['pos_cel_cov']
            src_dict['pos_cel_corr'] = o['pos_cel_corr']

            self.logger.info('Found source\n' +
                             'name: %s\n' % name +
                             'ts: %f' % p['amp'] ** 2)

            names.append(name)
            src_dicts.append(src_dict)

        return names, src_dicts

    def _find_sources_iterate(self, prefix, iiter, **kwargs):

        src_dict_template = kwargs.pop('model')

        threshold = kwargs.get('sqrt_ts_threshold')
        multithread = kwargs.get('multithread', False)
        min_separation = kwargs.get('min_separation')
        sources_per_iter = kwargs.get('sources_per_iter')
        search_skydir = kwargs.get('search_skydir', None)
        search_minmax_radius = kwargs.get('search_minmax_radius', [None, 1.0])
        tsmap_fitter = kwargs.get('tsmap_fitter', 'tsmap')
        free_params = kwargs.get('free_params', None)
        if not free_params:
            free_params = None

        if tsmap_fitter == 'tsmap':
            kw = kwargs.get('tsmap', {})
            kw['model'] = src_dict_template
            kw['multithread'] = multithread
            m = self.tsmap(utils.join_strings([prefix,
                                               'sourcefind_%02i' % iiter]),
                           **kw)

        elif tsmap_fitter == 'tscube':
            kw = kwargs.get('tscube', {})
            kw['model'] = src_dict_template
            kw['do_sed'] = False
            m = self.tscube(utils.join_strings([prefix,
                                                'sourcefind_%02i' % iiter]),
                            **kw)
        else:
            raise Exception(
                'Unrecognized option for fitter: %s.' % tsmap_fitter)

        if tsmap_fitter == 'tsmap':
            peaks = find_peaks(m['sqrt_ts'], threshold, min_separation)
            (names, src_dicts) = \
                self._build_src_dicts_from_peaks(peaks, m, src_dict_template)
        elif tsmap_fitter == 'tscube':
            peaks = find_peaks(m['sqrt_ts'], threshold, min_separation)
            (names, src_dicts) = \
                self._build_src_dicts_from_peaks(peaks, m, src_dict_template)
            """
            sd = m['tscube'].find_sources(threshold ** 2, min_separation,
                                          use_cumul=False,
                                          output_src_dicts=True,
                                          output_peaks=True)
            peaks = sd['Peaks']
            names = sd['Names']
            src_dicts = sd['SrcDicts']
            """

        # Loop over the seeds and add them to the model
        new_src_names = []
        for name, src_dict in zip(names, src_dicts):
            # Protect against finding the same source twice
            if self.roi.has_source(name):
                self.logger.info('Source %s found again.  Ignoring it.' % name)
                continue
            # Skip the source if it's outside the search region
            if search_skydir is not None:

                skydir = SkyCoord(src_dict['ra'], src_dict['dec'], unit='deg')
                separation = search_skydir.separation(skydir).deg

                if not utils.apply_minmax_selection(separation,
                                                    search_minmax_radius):
                    self.logger.info('Source %s outside of '
                                     'search region.  Ignoring it.',
                                     name)
                    continue

            self.add_source(name, src_dict, free=True)
            self.free_source(name, False)
            new_src_names.append(name)

            if len(new_src_names) >= sources_per_iter:
                break

        # Re-fit spectral parameters of each source individually
        for name in new_src_names:
            self.logger.info('Performing spectral fit for %s.', name)
            self.logger.debug(pprint.pformat(self.roi[name].params))
            self.free_source(name, True, pars=free_params)
            self.fit()
            self.logger.info(pprint.pformat(self.roi[name].params))
            self.free_source(name, False)

        srcs = []
        for name in new_src_names:
            srcs.append(self.roi[name])

        return srcs, peaks

    def localize(self, name, **kwargs):
        """Find the best-fit position of a source.  Localization is
        performed in two steps.  First a TS map is computed centered
        on the source with half-width set by ``dtheta_max``.  A fit is
        then performed to the maximum TS peak in this map.  The source
        position is then further refined by scanning the likelihood in
        the vicinity of the peak found in the first step.  The size of
        the scan region is set to encompass the 99% positional
        uncertainty contour as determined from the peak fit.

        Parameters
        ----------
        name : str
            Source name.

        {options}

        optimizer : dict
            Dictionary that overrides the default optimizer settings.

        Returns
        -------
        localize : dict
            Dictionary containing results of the localization
            analysis.

        """
        timer = Timer.create(start=True)
        name = self.roi.get_source_by_name(name).name

        schema = ConfigSchema(self.defaults['localize'],
                              optimizer=self.defaults['optimizer'])
        schema.add_option('use_cache', True)
        schema.add_option('prefix', '')
        config = utils.create_dict(self.config['localize'],
                                   optimizer=self.config['optimizer'])
        config = schema.create_config(config, **kwargs)

        self.logger.info('Running localization for %s' % name)

        free_state = FreeParameterState(self)
        loc = self._localize(name, **config)
        free_state.restore()

        self.logger.info('Finished localization.')

        if config['make_plots']:
            self._plotter.make_localization_plots(loc, self.roi,
                                                  prefix=config['prefix'])

        outfile = \
            utils.format_filename(self.workdir, 'loc',
                                  prefix=[config['prefix'],
                                          name.lower().replace(' ', '_')])

        if config['write_fits']:
            loc['file'] = os.path.basename(outfile) + '.fits'
            self._make_localize_fits(loc, outfile + '.fits',
                                     **config)

        if config['write_npy']:
            np.save(outfile + '.npy', dict(loc))

        self.logger.info('Execution time: %.2f s', timer.elapsed_time)
        return loc

    def _make_localize_fits(self, loc, filename, **kwargs):

        tab = fits_utils.dict_to_table(loc)
        hdu_data = fits.table_to_hdu(tab)
        hdu_data.name = 'LOC_DATA'

        hdus = [loc['tsmap_peak'].make_hdu(hdu='PRIMARY'),
                loc['tsmap'].make_hdu(hdu='TSMAP'),
                hdu_data]

        hdus[0].header['CONFIG'] = json.dumps(loc['config'])
        hdus[2].header['CONFIG'] = json.dumps(loc['config'])
        fits_utils.write_hdus(hdus, filename)

    def _localize(self, name, **kwargs):

        nstep = kwargs.get('nstep')
        dtheta_max = kwargs.get('dtheta_max')
        update = kwargs.get('update', True)
        prefix = kwargs.get('prefix', '')
        use_cache = kwargs.get('use_cache', False)
        free_background = kwargs.get('free_background', False)
        free_radius = kwargs.get('free_radius', None)
        fix_shape = kwargs.get('fix_shape', False)
        tsmap_fitter = kwargs.get('tsmap_fitter', 'tsmap')

        saved_state = LikelihoodState(self.like)
        loglike_init = -self.like()
        self.logger.debug('Initial Model Log-Likelihood: %f', loglike_init)

        if not free_background:
            self.free_sources(free=False, loglevel=logging.DEBUG)

        if free_radius is not None:
            diff_sources = [s.name for s in self.roi.sources if s.diffuse]
            skydir = self.roi[name].skydir
            free_srcs = [s.name for s in
                         self.roi.get_sources(skydir=skydir,
                                              distance=free_radius,
                                              exclude=diff_sources)]
            self.free_sources_by_name(free_srcs, pars='norm',
                                      loglevel=logging.DEBUG)

        src = self.roi.copy_source(name)
        skydir = src.skydir
        skywcs = self.geom.wcs
        src_pix = skydir.to_pixel(skywcs)

        fit0 = self._fit_position_tsmap(name, prefix=prefix,
                                        dtheta_max=dtheta_max,
                                        zmin=-3.0,
                                        use_pylike=False,
                                        tsmap_fitter=tsmap_fitter)

        self.logger.debug('Completed localization with TS Map.\n'
                          '(ra,dec) = (%10.4f,%10.4f) '
                          '(glon,glat) = (%10.4f,%10.4f)',
                          fit0['ra'], fit0['dec'],
                          fit0['glon'], fit0['glat'])

        # Fit baseline (point-source) model
        self.free_source(name, loglevel=logging.DEBUG)
        if fix_shape:
            self.free_source(name, free=False, pars='shape',
                             loglevel=logging.DEBUG)
        fit_output = self._fit(loglevel=logging.DEBUG, **
                               kwargs.get('optimizer', {}))

        # Save likelihood value for baseline fit
        loglike_base = fit_output['loglike']
        self.logger.debug('Baseline Model Log-Likelihood: %f', loglike_base)

        o = defaults.make_default_tuple(defaults.localize_output)
        o.name = name
        o.config = kwargs
        o.fit_success = True
        o.loglike_init = loglike_init
        o.loglike_base = loglike_base
        o.loglike_loc = np.nan
        o.dloglike_loc = np.nan

        if fit0['fit_success']:
            scan_cdelt = 2.0 * fit0['pos_r95'] / (nstep - 1.0)
        else:
            scan_cdelt = np.abs(skywcs.wcs.cdelt[0])

        self.logger.debug('Refining localization search to '
                          'region of width: %.4f deg',
                          scan_cdelt * nstep)

        fit1 = self._fit_position_scan(name,
                                       skydir=fit0['skydir'],
                                       scan_cdelt=scan_cdelt,
                                       **kwargs)

        o.loglike_loc = fit1['loglike']
        o.dloglike_loc = o.loglike_loc - o.loglike_base
        o.tsmap = fit0.pop('tsmap')
        o.tsmap_peak = fit1.pop('tsmap')
        # o.update(fit1)

        # Best fit position and uncertainty from fit to TS map
        o.fit_init = fit0

        # Best fit position and uncertainty from pylike scan
        o.fit_scan = fit1
        o.update(fit1)

        cdelt0 = np.abs(skywcs.wcs.cdelt[0])
        cdelt1 = np.abs(skywcs.wcs.cdelt[1])
        pix = fit1['skydir'].to_pixel(skywcs)
        o.pos_offset = skydir.separation(fit1['skydir']).deg
        o.xpix = float(pix[0])
        o.ypix = float(pix[1])
        o.deltax = (o.xpix - src_pix[0]) * cdelt0
        o.deltay = (o.ypix - src_pix[1]) * cdelt1

        o.ra_preloc = skydir.ra.deg
        o.dec_preloc = skydir.dec.deg
        o.glon_preloc = skydir.galactic.l.deg
        o.glat_preloc = skydir.galactic.b.deg

        if o.pos_offset > dtheta_max:
            o.fit_success = False

        if not o.fit_success:
            self.logger.warning('Fit to localization contour failed.')
        elif not o.fit_inbounds:
            self.logger.warning('Best-fit position outside of search region.')
        else:
            self.logger.info('Localization succeeded.')

        if update and ((not o.fit_success) or (not o.fit_inbounds)):
            self.logger.warning(
                'Localization failed.  Keeping existing position.')

        if update and o.fit_success and o.fit_inbounds:
            self.logger.info('Updating source %s '
                             'to localized position.', name)
            src = self.delete_source(name)
            src.set_position(fit1['skydir'])
            self.add_source(name, src, free=True)
            self.free_source(name, loglevel=logging.DEBUG)
            if fix_shape:
                self.free_source(name, free=False, pars='shape',
                                 loglevel=logging.DEBUG)

            fit_output = self.fit(loglevel=logging.DEBUG)
            o.loglike_loc = fit_output['loglike']
            o.dloglike_loc = o.loglike_loc - o.loglike_base
            src = self.roi.get_source_by_name(name)

            src['glon_err'] = o.glon_err
            src['glat_err'] = o.glat_err
            src['ra_err'] = o.ra_err
            src['dec_err'] = o.dec_err
            src['pos_err'] = o.pos_err
            src['pos_err_semimajor'] = o.pos_err_semimajor
            src['pos_err_semiminor'] = o.pos_err_semiminor
            src['pos_r68'] = o.pos_r68
            src['pos_r95'] = o.pos_r95
            src['pos_r99'] = o.pos_r99
            src['pos_angle'] = o.pos_angle
            src['pos_gal_cov'] = o.pos_gal_cov
            src['pos_gal_corr'] = o.pos_gal_corr
            src['pos_cel_cov'] = o.pos_cel_cov
            src['pos_cel_corr'] = o.pos_cel_corr
        else:
            saved_state.restore()
            self._sync_params(name)
            self._update_roi()

        self.logger.info('Localization completed with new position:\n'
                         '(  ra, dec) = (%10.4f +/- %8.4f,%10.4f +/- %8.4f)\n'
                         '(glon,glat) = (%10.4f +/- %8.4f,%10.4f +/- %8.4f)\n'
                         'offset = %8.4f r68 = %8.4f r95 = %8.4f r99 = %8.4f',
                         o.ra, o.ra_err, o.dec, o.dec_err,
                         o.glon, o.glon_err, o.glat, o.glat_err,
                         o.pos_offset, o.pos_r68, o.pos_r95, o.pos_r99)
        self.logger.info('LogLike: %12.3f DeltaLogLike: %12.3f',
                         o.loglike_loc, o.loglike_loc - o.loglike_init)

        return o

    def _fit_position(self, name, **kwargs):

        dtheta_max = kwargs.setdefault('dtheta_max', 0.5)
        nstep = kwargs.setdefault('nstep', 5)
        fit0 = self._fit_position_tsmap(name, **kwargs)

        if np.isfinite(fit0['pos_r68']):
            scan_cdelt = min(2.0 * fit0['pos_r68'] / (nstep - 1.0),
                             self._binsz)
        else:
            scan_cdelt = self._binsz

        fit1 = self._fit_position_scan(name,
                                       skydir=fit0['skydir'],
                                       scan_cdelt=scan_cdelt,
                                       **kwargs)
        return fit1, fit0

    def _fit_position_tsmap(self, name, **kwargs):
        """Localize a source from its TS map."""

        prefix = kwargs.get('prefix', '')
        dtheta_max = kwargs.get('dtheta_max', 0.5)
        zmin = kwargs.get('zmin', -3.0)
        tsmap_fitter = kwargs.get('tsmap_fitter', 'tsmap')

        kw = {'map_size': 2.0 * dtheta_max,
              'write_fits':  kwargs.get('write_fits', False),
              'write_npy':  kwargs.get('write_npy', False),
              'max_kernel_radius': self.config['tsmap']['max_kernel_radius'],
              'loglevel': logging.DEBUG}

        src = self.roi.copy_source(name)

        if src['SpatialModel'] in ['RadialDisk', 'RadialGaussian']:
            kw['max_kernel_radius'] = max(kw['max_kernel_radius'],
                                          2.0 * src['SpatialWidth'])

        skydir = kwargs.get('skydir', src.skydir)

        if tsmap_fitter == 'tsmap':
            tsmap = self.tsmap(utils.join_strings([prefix, name.lower().
                                                   replace(' ', '_')]),
                               model=src.data,
                               map_skydir=skydir,
                               exclude=[name],
                               use_pylike=kwargs.get('use_pylike', True),
                               make_plots=False, **kw)
        else:
            tsmap = self.tscube(utils.join_strings([prefix, name.lower().replace(' ', '_')]),
                                model=src.data,
                                map_skydir=skydir,
                                exclude=[name],
                                make_plots=False,
                                do_sed=False,
                                **kw)

        # Find peaks with TS > 4
        peaks = find_peaks(tsmap['ts'], 4.0, 0.2)
        peak_best = None
        o = {}
        for p in sorted(peaks, key=lambda t: t['amp'], reverse=True):
            xy = p['ix'], p['iy']
            ts_value = tsmap['ts'].data[xy[1], xy[0]]
            posfit = fit_error_ellipse(tsmap['ts'], xy=xy, dpix=2,
                                       zmin=max(zmin, -ts_value * 0.5))
            offset = posfit['skydir'].separation(self.roi[name].skydir).deg
            if posfit['fit_success'] and posfit['fit_inbounds']:
                peak_best = p
                break

        if peak_best is None:
            ts_value = np.max(tsmap['ts'].data)
            posfit = fit_error_ellipse(tsmap['ts'], dpix=2,
                                       zmin=max(zmin, -ts_value * 0.5))

        o.update(posfit)
        pix = posfit['skydir'].to_pixel(self.geom.wcs)
        o['xpix'] = float(pix[0])
        o['ypix'] = float(pix[1])
        o['skydir'] = posfit['skydir'].transform_to('icrs')
        o['pos_offset'] = posfit['skydir'].separation(
            self.roi[name].skydir).deg
        o['loglike'] = 0.5 * posfit['zoffset']
        o['tsmap'] = tsmap['ts']

        return o

    def _fit_position_scan(self, name, **kwargs):

        zmin = kwargs.get('zmin', -9.0)
        tsmap, loglike = self._scan_position(name, **kwargs)
        ts_value = np.max(tsmap.data)
        posfit = fit_error_ellipse(tsmap, dpix=2,
                                   zmin=max(zmin, -ts_value * 0.5))
        pix = posfit['skydir'].to_pixel(self.geom.wcs)

        o = {}
        o.update(posfit)
        o['xpix'] = float(pix[0])
        o['ypix'] = float(pix[1])
        o['skydir'] = posfit['skydir'].transform_to('icrs')
        o['pos_offset'] = posfit['skydir'].separation(
            self.roi[name].skydir).deg
        o['loglike'] = 0.5 * posfit['zoffset'] + loglike
        o['tsmap'] = tsmap

        return o

    def _scan_position(self, name, **kwargs):

        saved_state = LikelihoodState(self.like)

        skydir = kwargs.pop('skydir', self.roi[name].skydir)
        scan_cdelt = kwargs.pop('scan_cdelt', 0.02)
        nstep = kwargs.pop('nstep', 5)
        use_cache = kwargs.get('use_cache', True)
        use_pylike = kwargs.get('use_pylike', False)
        optimizer = kwargs.get('optimizer', {})

        # Fit without source
        self.zero_source(name, loglevel=logging.DEBUG)
        fit_output_nosrc = self._fit(loglevel=logging.DEBUG,
                                     **optimizer)
        self.unzero_source(name, loglevel=logging.DEBUG)
        saved_state.restore()
        self.free_norm(name, loglevel=logging.DEBUG)

        lnlmap = WcsNDMap.create(skydir=skydir, binsz=scan_cdelt, npix=(nstep, nstep),
                                 coordsys=wcs_utils.get_coordsys(self.geom.wcs))

        src = self.roi.copy_source(name)

        if use_cache and not use_pylike:
            self._create_srcmap_cache(src.name, src)

        coord = MapCoord.create(lnlmap.geom.get_coord(flat=True),
                                coordsys=lnlmap.geom.coordsys)
        scan_skydir = coord.skycoord.icrs
        for lon, lat, ra, dec in zip(coord.lon, coord.lat,
                                     scan_skydir.ra.deg, scan_skydir.dec.deg):

            spatial_pars = {'ra': ra, 'dec': dec}
            self.set_source_morphology(name,
                                       spatial_pars=spatial_pars,
                                       use_pylike=use_pylike)
            fit_output = self._fit(loglevel=logging.DEBUG,
                                   **optimizer)
            lnlmap.set_by_coord((lon, lat), fit_output['loglike'])

        self.set_source_morphology(name, spatial_pars=src.spatial_pars,
                                   use_pylike=use_pylike)
        saved_state.restore()

        lnlmap.data -= fit_output_nosrc['loglike']
        tsmap = WcsNDMap(lnlmap.geom, 2.0 * lnlmap.data)

        self._clear_srcmap_cache()
        return tsmap, fit_output_nosrc['loglike']

    def _fit_position_opt(self, name, use_cache=True):

        state = SourceMapState(self.like, [name])

        src = self.roi.copy_source(name)

        if use_cache:
            self._create_srcmap_cache(src.name, src)

        loglike = []
        skydir = src.skydir
        skywcs = self.geom.wcs
        src_pix = skydir.to_pixel(skywcs)

        c = skydir.transform_to('icrs')
        src.set_radec(c.ra.deg, c.dec.deg)
        self._update_srcmap(src.name, src)

        print(src_pix, self.like())

        import time

        def fit_fn(params):

            t0 = time.time()

            c = SkyCoord.from_pixel(params[0], params[1], self.geom.wcs)
            c = c.transform_to('icrs')
            src.set_radec(c.ra.deg, c.dec.deg)

            t1 = time.time()

            self._update_srcmap(src.name, src)

            t2 = time.time()

            val = self.like()

            t3 = time.time()

            print(params, val)
            # print(t1-t0,t2-t1,t3-t2)

            return val

        #lnl0 = fit_fn(src_pix[0],src_pix[1])
        #lnl1 = fit_fn(src_pix[0]+0.1,src_pix[1])
        # print(lnl0,lnl1)

        import scipy

        #src_pix[1] += 3.0
        p0 = [src_pix[0], src_pix[1]]

        #p0 = np.array([14.665692574327048, 16.004594098101926])
        #delta = np.array([0.3,-0.4])
        #p0 = [14.665692574327048, 16.004594098101926]

        o = scipy.optimize.minimize(fit_fn, p0,
                                    bounds=[(0.0, 39.0),
                                            (0.0, 39.0)],
                                    # method='L-BFGS-B',
                                    method='SLSQP',
                                    tol=1e-6)

        print('fit 2')

        o = scipy.optimize.minimize(fit_fn, o.x,
                                    bounds=[(0.0, 39.0),
                                            (0.0, 39.0)],
                                    # method='L-BFGS-B',
                                    method='SLSQP',
                                    tol=1e-6)
        print(o)

        print(fit_fn(p0))
        print(fit_fn(o.x))
        print(fit_fn(o.x + np.array([0.02, 0.02])))
        print(fit_fn(o.x + np.array([0.02, -0.02])))
        print(fit_fn(o.x + np.array([-0.02, 0.02])))
        print(fit_fn(o.x + np.array([-0.02, -0.02])))

        state.restore()

        return o


# ZBFitter5.py
# Reverting to radial scaling factors UVMAX=27mrad, XYmax=390mm: RT169.py
# Reading RT169_Nrings_8_Rmax_27.txt
#
#    0.000023    0.000002   -0.013169   -0.000001    0.000012    0.018778    0.000187   -0.000000   -0.000005    0.000001    0.000000
#    0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000004    0.000001    0.000002    0.000005
#
#  all 24 look good!
#  Maybe a Diameter effect at the focal surface?
#  DESI-9.OPT has 840mm: should be plenty!
# Let's try eliminating that Diameter:  DESI-10.OPT go to  RT170.py and YES here we are:
#
# Here using radial scaling factors UVMAX=28mrad,  XYMAX=407mm
# Reading RT170_Nrings_8_Rmax_28.txt  using DESI-10.OPT lacking FP Diameter
#
#  Results forward, 11 ZB coefs, RT170, ZBF5fwd, UVMAX=28
# ADC0      Xtrans      Ytrans      Magnif    UpDownUp  RightLeftRight  InOut      OutInOut     Roll   LeftRightCurl  UpDownCurl  RollAntiRoll
# coefs    0.000023    0.000001   -0.019117   -0.000001    0.000012    0.020128    0.000215   -0.000000   -0.000005    0.000001    0.000000
# sigma    0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000000    0.000001    0.000002    0.000005
# ADC6 (180 deg):
# coefs   -0.012770    0.000001   -0.019101   -0.000001   -0.000699    0.020127    0.000214    0.000000   -0.000104    0.000001   -0.000000
# sigma    0.000001    0.000001    0.000001    0.000002    0.000002    0.000002    0.000003    0.000003    0.000002    0.000002    0.000005
#     .. in fact all the 25 ADC settings look good. No "nan" or crazy results.
#
#  Results reverse, 11 ZB coefs, RT170, ZBF5fwd, UVMAX=28
# ADC0      Xtrans      Ytrans      Magnif    UpDownUp  RightLeftRight  InOut      OutInOut     Roll   LeftRightCurl  UpDownCurl  RollAntiRoll
# coefs   -0.000023   -0.000001    0.018720    0.000001   -0.000010   -0.020395    0.000566    0.000000    0.000005   -0.000001    0.000000
# sigma    0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000000    0.000001    0.000002    0.000005
#     .. in fact ALL the 25 ADC settings look good.  No "nan" or crazy results.
#
# Conclusion: use DESI-10.OPT prescription with its absence of FP Diameter spec.
#
#
# Earlier:  using radial scaling factors UVMAX=28mrad,  XYMAX=407mm
# Reading RT168_Nrings_8_Rmax_28.txt
#
#  Results forward, 11 ZB coefs, RT168, ZBF3fwd, UVMAX=28
# ADC0      Xtrans      Ytrans      Magnif    UpDownUp  RightLeftRight  InOut      OutInOut     Roll   LeftRightCurl  UpDownCurl  RollAntiRoll
# coefs    0.000023   -0.000000   -0.019118    0.000000    0.000012    0.020126    0.000236   -0.000000   -0.000005    0.000000    0.000000
# sigma    0.000001    0.000005    0.000001         nan    0.000001    0.000002    0.000002    0.000007    0.000001         nan    0.000008
# ADC6 (180 deg)
# coefs   -0.012771   -0.000000   -0.019103    0.000000   -0.000701    0.020126    0.000236   -0.000000   -0.000108    0.000000    0.000000
# sigma    0.000001    0.000006    0.000002  631.809570    0.000002    0.000002    0.000003    0.000010    0.000002  631.809566    0.000010
#                                             ^^^^^^^^                                                                ^^^^^^^^
#  Yikes, why do the sigmas go crazy at ADC settings of 0, 180 degrees?
#  Could it be my UVMAX=28, or is it some other bug?
#  To find out, see RT169: just like RT168 but with UVMAX=27.
#
# Results reverse, 11 ZB coefs, RT168, ZBF3rev,UVMAX=28
# ADC0:
# coefs    -0.000023   -0.000000    0.018720   -0.000000   -0.000010   -0.020398    0.000538    0.000000    0.000005   -0.000000   -0.000000
# sigma     0.000001    0.000005    0.000001   52.833776    0.000001    0.000002    0.000002    0.000008    0.000001   52.833778    0.000008
# ADC6:
# coefs     0.012763    0.000000    0.018716   -0.000000   -0.000237   -0.020398    0.000537    0.000000    0.000271   -0.000000   -0.000000
# sigma     0.000001    0.000005    0.000001  348.925720    0.000001    0.000002    0.000002    0.000009    0.000001  348.925720    0.000009
#
#  COMPARE PREVIOUS RESULTS USING ZBFitter2.py:
#
# Results forward, 11 ZB coefs, RT166, ZBF2fwd, UVMAX=27
# ADC0      Xtrans      Ytrans      Magnif    UpDownUp  RightLeftRight  InOut      OutInOut     Roll   LeftRightCurl  UpDownCurl  RollAntiRoll
# coefs     0.000023    0.000002   -0.013169   -0.000001    0.000012    0.018778    0.000187   -0.000000   -0.000005    0.000001    0.000000
# sigma     0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000004    0.000001    0.000002    0.000005
# ADC6:
# coefs    -0.013285    0.000002   -0.013153   -0.000001   -0.000678    0.018777    0.000187    0.000000   -0.000101    0.000001   -0.000000
# sigma     0.000001    0.000001    0.000002    0.000002    0.000001    0.000002    0.000002    0.000004    0.000001    0.000002    0.000004
#
# Results reverse, 11 ZB coefs, RT166, ZBF2rev,UVMAX=27
# ADC0:
# coefs    -0.000023   -0.000002    0.012965    0.000001   -0.000010   -0.018778    0.000475    0.000000    0.000005   -0.000001    0.000000
# sigma     0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000000    0.000001    0.000002    0.000005
# ADC6:
# coefs     0.013233   -0.000002    0.012961    0.000001   -0.000226   -0.018777    0.000473   -0.000000    0.000256   -0.000001   -0.000000
# sigma     0.000001    0.000001    0.000001    0.000002    0.000001    0.000002    0.000002    0.000001    0.000001    0.000002    0.000000

# coefs  -0.000023   -0.000002    0.012965    0.000001   -0.000010   -0.018778    0.000475    0.000000    0.000005   -0.000001    0.000000
#
# Modelling ADC roll as Zhao-Burge distortions.
# Assumes that RT168.py has already run a complete monochromatic trace,
# with ~ 400 rays per pupil, and 217 sky locations with image centroids,
# for ADC1 = 0,30,60,... and then ADC2 = 0, 30, 60...degrees
# These 24 cases are built into RT166.
# The results file with all 24 cases is RT166_Nrings_8_Rmax_27.txt
#

import csv
import numpy as np
from numpy import genfromtxt
import matplotlib.pyplot as plt
import math
from scipy.optimize import curve_fit

PROGNAME = 'ZBFitter5.py'

LUT = [0,  1,  2,  5,  6,   9,   20,  27, 28, 29, 30]   # 11 polynomials
# parm 0,  1,  2,  3,  4,   5,   6,   7,  8,  9,  10    # parm numbers
# ZB:  S2, S3, S4, S7, S8,  S11, S22, T4, T7, T8, T11    # Zhao-Burge labels
#      Xt, Yt, mag,

NPARMS = len(LUT)

infilename = 'RT169_Nrings_8_Rmax_27.txt'  # 5425 records = 25 x 217

def getFileArray(filename):
    nums2D = list()
    with open(filename) as f:
        data = f.readlines()
    for row in data:
        numrow = list()
        words = row.split()
        for word in words:
            try:
                x = float(word)
            except:
                x = -0.0
            numrow.append(x)
        nums2D.append(numrow)
    #---Make nums2D rectangular, lest asarray() will return a 2D list---
    nrows = len(nums2D)
    ncols = 0
    for i in range(nrows):
        ncols = max(ncols, len(nums2D[i]))
    print(' ncols = ' + str(ncols))
    for i in range(nrows):
        while len(nums2D[i]) < ncols:
            nums2D[i].append(-0.0)
    myArray = np.asarray(nums2D)
    return myArray

#-------ZERNIKE NORMALIZED FUNCTIONS & DERIVATIVES--------
#-------------Using {n,m} definitions---------------------

def factorial(n):
    if n > 1:
       return int(n*factorial(n-1))
    else:
       return 1

def convertNolltoBW(noll):
    # converts a Noll Zernike index to the B&W {n,m,t} triplet
    n = int(-0.5+np.sqrt(2*noll-1.75))
    m = 0
    base = int(0.5*n**2 +0.5*n + 1)
    diff = noll - base
    if n%2==0:
        m = 2*int(0.5*diff + 0.7)
    else:
        m = 2*int(0.5*diff + 1.2) - 1
    if noll%2>0:
        m = -m
    return np.array([n, m])

def convertWyanttoBW(wyant):
    # converts a Wyant Zernike index to the B&W {n,m,t}
    halfsum = int(np.sqrt(wyant))
    idif = int(wyant - halfsum**2)
    halfdif = int (idif/2)
    n = halfsum + halfdif
    m = halfsum - halfdif
    if idif%2 >0:
        m = -m
    return np.array([n, m])

def getZernFuncXY(nm, xnorm, ynorm):   # BIG NINE:  #1
    # Here, xnorm and ynorm must lie within the unit circle
    rnorm = np.sqrt(xnorm*xnorm + ynorm*ynorm)
    angle = np.arctan2(ynorm,xnorm)
    return getZernRadial(nm,rnorm) * getZernAngular(nm,angle)

def getZernRadial(nm, rnorm):    # BIG NINE: #4
    n = nm[0]             # B&W
    m = np.abs(nm[1])     # B&W
    halfsum = (n+m)/2
    idif = n-m
    halfdif = int(idif/2)
    # n = halfsum + halfdif   # or, halfsum = (n+m)/2
    # m = halfsum - halfdif   # or, halfdif = (n-m)/2
    # loop through the polynomial
    result = 0.
    for i in range(0, halfdif+1):
        expon = int(n-2*i)
        sign = 1 if i%2 == 0 else -1
        numer = sign * factorial(n-i)
        denom = factorial(i) * factorial(halfsum-i) * factorial(halfdif-i)
        coef = numer / denom
        term = coef*math.pow(rnorm, expon)
        result = result + term
    return result

def getZernAngular(nm, theta):
    m = nm[1]    # B&W
    if m==0:
        return 1.
    if m>0:
        return math.cos(m*theta)
    m = np.abs(m)         # note this abs() function
    return math.sin(m*theta)

def zernFormulaText(nm):           # BIG NINE: #8
    #---generates a text representation of the specified Zernike function--
    n = nm[0]  # B&W
    m = nm[1]  # B&W
    # print 'New zernFormulaText() is using n, m: ', n, m
    needsine = True if m<0 else False
    m = np.abs(m)
    halfsum = (n+m)/2
    idif = n-m
    halfdif = int(idif/2)

    #--first do the radial part----
    nterms = 0
    s = ''
    #--evaluate the radial polynomial-----
    for i in range(0, halfdif+1):
        nterms = nterms + 1
        # print "Starting with n,  m, i, nterms = ", n, m, i, nterms
        expon = int(n-2*i)                # start with highest exponent
        # print "  expon = ", expon
        sign = 1 if i%2 == 0 else -1      # alternating signs in Zernike series
        strsign = '+' if i%2==0 else '-'
        numer = sign * factorial(n-i)
        denom = factorial(i) * factorial(halfsum-i) * factorial(halfdif-i)
        coef = numer / denom
        scoef = str(coef)
        if coef==1 and expon>0:          # suppress showing coef=1
            scoef = ''
        if coef > 0 and nterms > 1:
            scoef = '+' + scoef
        s = s + scoef
        if expon > 0:
            s = s + 'r'
        if expon > 1:
            s = s + '^' + str(expon)
    if nterms>1 and m!=0:
        s = '('+ s + ')'
    #--then do the azimuthal part, if any--------
    if m==0:
        return s
    strm = ''
    if m>1:
        strm = str(m)
    if needsine:
        s = s + '*sin(' + strm + 't)'
    else:
        s = s + '*cos(' + strm + 't)'
    return s

#--END ZERNIKES with {n,m} B&W indexing------------

#-----Zhao-Burge functions built from Zernikes---------------------

rh = np.sqrt(0.5)
rt = np.sqrt(2.0)

NPARMS = len(LUT)

def getZ(noll, x, y):
    return getZernFuncXY(convertNolltoBW(noll), x, y)

def getZhaoBurgeTerm(whichparm, x, y):
    # Given a Lampton index "which" 0....19 and an object point x, y,
    # fetches the needed the Zernikes via their Noll numbers..
    # Returns the modeled image point deviation x, y.
    # Case numbers shown are from Zhao & Burge Tables 1 and 2.
    which = LUT[whichparm]
    if which==0:   # case "S2"  r^0, keep; X translate
        return getZ(1,x,y), 0.0

    if which==1:   # case "S3"  r^0, keep; Y translate
        return 0.0, getZ(1,x,y)

    if which==2:   # case "S4"   r^1, keep: magnification
        return rh*getZ(2,x,y), rh*getZ(3,x,y)

    if which==3:   # case "S5"   r^1   -1ppm Mangled, 3ppm PartlyMangled, 4ppm ADC45
        return rh*getZ(3,x,y), rh*getZ(2,x,y)

    if which==4:   # case "S6"    r^1  24ppm Mangled,  6ppm PartlyMangled,-5ppm ADC45
        return rh*getZ(2,x,y), -rh*getZ(3,x,y)

    if which==5:   # case "S7"   r^2 125ppm Mangled  -908ppm PartlyMangled, ADC=423
        return 0.5*getZ(5,x,y), rh*getZ(4,x,y)-0.5*getZ(6,x,y)

    if which==6:   # case "S8"    r^2, 1223 Mangled, 653 PartlyMangled 10ppm ADC
        return rh*getZ(4,x,y)+0.5*getZ(6,x,y), 0.5*getZ(5,x,y)

    if which==7:   # case "S9"   r^3, 77ppm Mangled, 383 PartlyMangled, 116ppm ADC
        return rh*getZ(5,x,y), rh*getZ(6,x,y)

    if which==8:   # case "S10"   r^2, -6ppm Mangled  -1ppm PartlyMangled, zero ADC
        return rh*getZ(6,x,y), -rh*getZ(5,x,y)

    if which==9:   # case "S11"   r^3, huge
        return rh*getZ(8,x,y), rh*getZ(7,x,y)

    if which==10:  # case "S12"   r^3, 2ppm Mangled,zero PartlyMangled, zero ADC
        return 0.5*getZ(8,x,y)+0.5*getZ(10,x,y), -0.5*getZ(7,x,y)+0.5*getZ(9,x,y)

    if which==11:  # case "S13"    r^3, zero Mangled  zero PartlyMangled, 1ppm ADC
        return 0.5*getZ(7,x,y)+0.5*getZ(9,x,y), 0.5*getZ(8,x,y)-0.5*getZ(10,x,y)

    if which==12:  # case "S14"    r^3, 1ppm Mangled  zero PartlyMangled zero ppm ADC
        return rh*getZ(10,x,y), -rh*getZ(9,x,y)

    if which==13:  # case "S15"    r^3, zero Mangled   zeroPartlyMangled  zero ppm ADC
        return rh*getZ(9,x,y), rh*getZ(10,x,y)

    if which==14:  # Case "S16"  r^4  38ppm Mangled; 9ppm PartlyMangled, zeroADCero ADC
        return rh*getZ(11,x,y)+0.5*getZ(12,x,y), 0.5*getZ(3,x,y)

    if which==15:  # Case "S17":   r^4   15ppmMangled, -7ppm PartlyMangled,   8ppm ADC
        return 0.5*getZ(3,x,y), rh*getZ(11,x,y)-0.5*getZ(12,x,y)

    if which==16:  # Case "S18"  r^4   -6ppm Mangled, zero PartlyMangled, zero ADC
        return 0.5*getZ(12,x,y)+0.5*getZ(14,x,y), 0.5*getZ(15,x,y)-0.5*getZ(13,x,y)

    if which==17:  # Case "S19"   r^4   2ppm Mangled  -1ppm PartlyMangled  1ppm ADC;
        return 0.5*getZ(13,x,y)+0.5*getZ(15,x,y), 0.5*getZ(12,x,y)-0.5*getZ(14,x,y)

    if which==18:  # Case "S20"  r^4   zero Mangled, zero PartlyMangled, zero ADC
        return rh*getZ(14,x,y), -rh*getZ(15,x,y)

    if which==19:  # Case "S21"  r^4   zero Mangled  1ppm PartlyMangled, zero ADC
        return rh*getZ(15,x,y), rh*getZ(14,x,y)

    if which==20:  # Case "S22"  r^5  171 ppm Mangled  172ppm PartlyMangled, 172ppm ADC
        return rh*getZ(16,x,y), rh*getZ(17,x,y)

    if which==21:  # Case "S23"  r^5  zero Mangled  1ppm PartlyMangled,  zero ADC
        return 0.5*getZ(17,x,y)+0.5*getZ(19,x,y), 0.5*getZ(16,x,y)-0.5*getZ(18,x,y)

    if which==22:  # Case "S24"   r^5  -1ppm Mangled  zeroPartlyMangled,zero ADC
        return 0.5*getZ(16,x,y)+0.5*getZ(18,x,y), -0.5*getZ(17,x,y)+0.5*getZ(19,x,y)

    if which==23:  # Case "S25"  r^5   zero Mangled  zero PartlyMangled zero ADC
        return 0.5*getZ(19,x,y)+0.5*getZ(21,x,y), 0.5*getZ(18,x,y)-0.5*getZ(20,x,y)

    if which==24:  # Case "S26"   r^5  zero ppm  zero PartlyMangled  zero ADC
        return 0.5*getZ(18,x,y)+0.5*getZ(20,x,y), -0.5*getZ(19,x,y)+0.5*getZ(21,x,y)

    if which==25:  # Case "S27"   r^5   zero ppm  zero, PartlyMangled; zero ADC
        return rh*getZ(21,x,y), rh*getZ(20,x,y)

    if which==26:  # Case "S28"  r^5  zero ppm  zero PartlyMangled, zero ADC
        return rh*getZ(20,x,y), -rh*getZ(21,x,y)

    if which==27:  #  case "T4"   r^1,  huge, -6ppm ADC.  Roll.
        return rh*getZ(3,x,y), -rh*getZ(2,x,y)

    if which==28:  # case "T7"    r^2,  -265ppm Mangled, 131 PartlyMangled, -4ppm ADC
        return rh*getZ(4,x,y)-0.5*getZ(6,x,y), -0.5*getZ(5,x,y)

    if which==29:  # case "T8"   r^2, -105ppm Mangled.  -544ppm PartlyMangled, -163 ADC
        return 0.5*getZ(5,x,y), -rh*getZ(4,x,y)+0.5*getZ(6,x,y)

    if which==30:  # case "T11"   r^3, 358ppm Mangled, 358 PartlyMangled  zero ADC
        return rh*getZ(7,x,y), -rh*getZ(8,x,y)

    if which==31:  # case "T12"   r^3,  -1ppm Mangled, zero PartlyMangled, zero ADC
        return -0.5*getZ(7,x,y)+0.5*getZ(9,x,y), -0.5*getZ(8,x,y)-0.5*getZ(10,x,y)

    if which==32:  # case "T13"  r^3, +1ppm Mangled,  zero PartlyMangled, zero ADC
        return 0.5*getZ(8,x,y)-0.5*getZ(10,x,y), -0.5*getZ(7,x,y)-0.5*getZ(9,x,y)

    print("ZhaoBurgeTerm() is exitting because which = ", which)
    quit()

"""
def doRemapping(myParms):
    # feeds each sky point {u,v} into getZhaoSum()
    # will be called by optimizer to try out a mix of parameters.
    for i in range(0, len(SkyGrid)):
        u = SkyGrid[i, 0]
        v = SkyGrid[i, 1]
        x, y = getZhaoSum(myParms, u, v)
        # print 'result:  i, x, y ={:6d}{:12.6f}{:12.6f}'.format(i, x, y)
        X[i] = x
        Y[i] = y
"""

def zbFunc(uv, *parms): # uv = 1D array input concatenate(u[],v[])
    # Model: any number of Zhao & Burge parameters
    # note: curve_fit() calls its given func(args, *parms)
    # print("anyFunc() hasreceived parms = ")
    # print(["{0:0.6f}".format(xx) for xx in parms])
    pq = np.zeros(len(uv))
    half = len(uv)//2
    for i in range(0, half):  # EACH OBJECT POINT NOT AN ARRAY
        u = uv[i]
        v = uv[i+half]
        x, y = getZhaoSum(u, v, *parms)  # always args then parms
        pq[i] = x
        pq[i+half] = y
    return pq

def getZhaoSum(u, v, *coefs):
    # This models the distortion of ONE SKY POINT onto ONE FOCAL POINT.
    # Sums over all terms.
    # coefs is the array of 20 coefficients
    # u and v are Cartesian sky coordinates mapped into unit radius circle.
    x = u   # initially undeviated
    y = v   # initially undeviated
    for index in range(0, len(coefs)):
        dx, dy = getZhaoBurgeTerm(index, u, v)
        x += dx * coefs[index]
        y += dy * coefs[index]
    return x, y

def radialFunc(uv, *parms):
    # 6 parms model: dx, dy, a1, a3, a5, a7
    # untested!
    pq = np.zeros(len(uv))
    half = len(uv)//2
    for i in range(0, half):
        u = uv[i]
        v = uv[i+half]
        r = np.sqrt(u*u+v*v)
        if r==0.:
            u = 1E-12
        t = np.arctan2(u,v)
        rpoly = (1.+p[2])*r + p[3]*r**3 + p[4]*r**5 + p[5]*r**7
        x = p[0] + np.cos(t)*rpoly
        y = p[1] + np.sin(t)*rpoly
        pq[i] = x
        pq[i+half] = y
    return pq

#-------------------Main program-------------------------

alldata = getFileArray(infilename)
print(alldata.shape)  # 5208 x 9; 24 fields of 217 spots each
print('  ADC1  ADC2    RU        RV      Ngood    Xave      Yave       Xrms      Yrms')
Nadcs     = 25
Neach     = 217
NRINGS    = 8
UVMAX     = 28     # milliradians
XYMAX     = 407    # millimeters

# first get the forward (sky to FP) fits to the 25 ADC tasks

outfilename = 'ZBF5fwd_Nrings_' + str(NRINGS) + '_Rmax_' + str(UVMAX) + '.txt'

with open(outfilename, 'w') as outfile:
    for adc in range(Nadcs):
        adc12 = np.copy(alldata[adc*Neach:(adc+1)*Neach, 0:2])    # rows,cols
        uv    = np.copy(alldata[adc*Neach:(adc+1)*Neach, 2:4])    # rows,cols
        xy    = np.copy(alldata[adc*Neach:(adc+1)*Neach, 5:7])    # rows,cols
        uv *= 1000./UVMAX   # for unit circle scaling
        xy *= 1./XYMAX     # for unit circle scaling
        print()

        pinitial = np.zeros(NPARMS)
        args     = np.concatenate([uv[:,0], uv[:,1]])
        goals    = np.concatenate([xy[:,0], xy[:,1]])
        popt, covar = curve_fit(zbFunc, args, goals, pinitial)
        #  note: curve_fit() calls its given func(args, *parms)
        print('\npopt: \n', popt)
        perror = np.sqrt(covar.diagonal())
        print('\nerrors: \n', perror)
        combo = np.concatenate((popt, perror))  # double parentheses
        result = ''
        for item in combo:
            result += '{:12.6f}'.format(item)
        outfile.writelines(result + '\n') # yes, plural even for one line

# then get the reverse (FP to sky) fits to the 25 ADC tasks

outfilename = 'ZBF5rev_Nrings_' + str(NRINGS) + '_Rmax_' + str(UVMAX) + '.txt'

with open(outfilename, 'w') as outfile:
    for adc in range(Nadcs):
        adc12 = np.copy(alldata[adc*Neach:(adc+1)*Neach, 0:2])    # rows,cols
        uv    = np.copy(alldata[adc*Neach:(adc+1)*Neach, 2:4])    # rows,cols
        xy    = np.copy(alldata[adc*Neach:(adc+1)*Neach, 5:7])    # rows,cols
        uv *= 1000./UVMAX   # for unit circle scaling
        xy *= 1./XYMAX     # for unit circle scaling
        print()

        pinitial = np.zeros(NPARMS)
        goals    = np.concatenate([uv[:,0], uv[:,1]])
        args     = np.concatenate([xy[:,0], xy[:,1]])
        popt, covar = curve_fit(zbFunc, args, goals, pinitial)
        #  note: curve_fit() calls its given func(args, *parms)
        print('\npopt: \n', popt)
        perror = np.sqrt(covar.diagonal())
        print('\nerrors: \n', perror)
        combo = np.concatenate((popt, perror))  # double parentheses
        result = ''
        for item in combo:
            result += '{:12.6f}'.format(item)
        outfile.writelines(result + '\n') # yes, plural even for one line


#  Copyright (c) 2019, CNRS-LAAS
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  * Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#
#  * Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import typing as ty

import cv2
import numpy as np
import skimage.draw
import skimage.measure

import fire_rs.geodata.geo_data
import fire_rs.rbf

class Perimeter:

    def __init__(self, wildfire: fire_rs.geodata.geo_data.GeoData, threshold: float,
                 layer: str = 'ignition', perimeter_background=np.inf):
        self._wildfire = wildfire
        self._threshold = threshold
        self._layer = layer
        self._empty_val = perimeter_background

        self._perimeter_array, self._cells, self._contour = _compute_perimeter(
            self._wildfire, self._threshold, layer=self._layer, empty_val=self._empty_val)

        if self._contour:
            for cont in self._contour:
                skimage.measure.grid_points_in_poly(self._wildfire.data.shape, cont)

        self._perimeter_geodata = None
        self._area_array = None

    @property
    def array(self) -> np.ndarray:
        return self._perimeter_array

    @property
    def cells(self) -> ty.MutableMapping[ty.Tuple[int, int], np.float64]:
        return self._cells

    @property
    def count(self) -> int:
        """Number of indepedent perimeters"""
        return len(self._contour)

    @property
    def geodata(self) -> fire_rs.geodata.geo_data.GeoData:
        if self._perimeter_geodata is None:
            self._perimeter_geodata = self._wildfire.clone(data_array=self._perimeter_array,
                                                           dtype=self._wildfire.data.dtype)
        return self._perimeter_geodata

    @property
    def area_array(self):
        """Mask of cells inside the perimeter"""
        if self._area_array is None:
            if self._contour:
                self._area_array = skimage.measure.grid_points_in_poly(self._wildfire.data.shape,
                                                                       self._contour[0])
                if len(self._contour) > 1:
                    for cont in self._contour[1:]:
                        self._area_array += skimage.measure.grid_points_in_poly(
                            self._wildfire.data.shape, cont)
        return self._area_array

_compute_perimeter_output_type = ty.Tuple[
    np.ndarray, ty.MutableMapping[ty.Tuple[int, int], np.float64], ty.List[np.ndarray]]

def _compute_perimeter(wildfire: fire_rs.geodata.geo_data.GeoData, threshold: float,
                       layer: str = 'ignition', empty_val=np.inf) -> _compute_perimeter_output_type:
    """Extract the perimeter given by threshold from a fire map."""

    array = np.ones(wildfire.data.shape, dtype=np.float64) * empty_val
    cells = {}

    contours = skimage.measure.find_contours(wildfire.data[layer], threshold)

    for contour in contours:
        try:
            rr, cc = skimage.draw.polygon_perimeter(contour[..., 0], contour[..., 1],
                                                    shape=wildfire.data.shape, clip=True)
            # Set perimeter in array format
            array[rr, cc] = wildfire.data[layer][rr, cc]
            # Set perimeter in dict format
            for r, c in zip(rr, cc):
                cells[r, c] = wildfire.data[layer][r, c]
        except IndexError as e:
            pass # Ignore contour if it contains NaN (polygon_perimeter throws IndexError)

    # for contour in contours:
    #     prev_edge = contour[0]
    #     for edge in contour[1:]:
    #         if np.isnan(prev_edge).any() or np.isnan(edge).any():
    #             continue
    #         rr, cc = skimage.draw.line(*np.asarray(prev_edge, dtype=int),
    #                                    *np.asarray(edge, dtype=int))
    #         # Set perimeter in array format
    #         array[rr, cc] = wildfire.data[layer][rr, cc]
    #         # Set perimeter in dict format
    #         for r, c in zip(rr, cc):
    #             cells[r, c] = wildfire.data[layer][r, c]
    #
    #         prev_edge = edge

    return array, cells, contours

def interpolate(x, y, z, shape, function='thin_plate') -> np.ndarray:
    """RBF interpolation"""
    # Wildland fire modeling with an Eulerian level set method and automated calibration
    # might give a clue of which kind of kernel function to use

    # default smooth=0 for interpolation
    interpolator = fire_rs.rbf.Rbf(x, y, z, function=function, smooth=0, cond=10 ** -5)

    xi = np.linspace(0, shape[0] - 1, shape[0])
    yi = np.linspace(0, shape[1] - 1, shape[1])
    meshgrid = np.meshgrid(xi, yi, indexing="ij")

    dense_array = interpolator(*[x.flatten() for x in meshgrid])

    return dense_array.reshape(shape[0], shape[1])

def rate_of_spread_map(firemap: fire_rs.geodata.geo_data.GeoData, layer="ignition",
                       output_layer="ros") -> fire_rs.geodata.geo_data.GeoData:
    """Compute the Rate of Spread from a wildfire_map"""
    gradient = firemap.clone(data_array=np.linalg.norm(np.gradient(firemap[layer]), axis=0),
                             dtype=[(output_layer, 'float64')])
    return gradient

class WildfireGraph:
    """Computation of the propagation graph and end of ignition from a wildfire map"""
    IGNITION_END = 'ignition_end'
    PROP_X = 'prop_x'
    PROP_Y = 'prop_y'
    PROP_DIR = 'prop_dir'

    def __init__(self, firemap: fire_rs.geodata.geo_data.GeoData, ignition_layer: str = 'ignition',
                 ignition_end_layer: str = IGNITION_END, prop_x_layer: str = PROP_X,
                 prop_y_layer: str = PROP_Y, prop_dir_layer: str = PROP_DIR):
        self._ignition_layer = ignition_layer
        self._ignition_end_layer = ignition_end_layer
        self._prop_x_layer = prop_x_layer
        self._prop_y_layer = prop_y_layer
        self._prop_dir_layer = prop_dir_layer
        self.geodata = firemap.clone(fill_value=0,
                                     dtype=[(ignition_layer, 'float64'),
                                            ('ignition_end', 'float64'),
                                            ('prop_x', 'int8'),
                                            ('prop_y', 'int8'),
                                            ('prop_dir', 'float64')])
        self.geodata.data[ignition_layer] = firemap.data[ignition_layer]
        grad = np.gradient(self.geodata.data[ignition_layer])
        self.geodata.data[prop_dir_layer] = np.arctan2(grad[1], grad[0])
        self.geodata.data[prop_x_layer] = np.array(
            np.round(np.cos(self.geodata.data[prop_dir_layer])), np.int)
        self.geodata.data[prop_y_layer] = np.array(
            np.round(np.sin(self.geodata.data[prop_dir_layer])), np.int)
        self.geodata.data[ignition_end_layer] = WildfireGraph._compute_traversal_end(
            self.geodata.data[ignition_layer])

    @staticmethod
    def _compute_propagation_direction(fire_array: np.array):
        def default_ignition(x, y, dx, dy):
            if x == 0 and dx < 0:
                return fire_array[x, y]
            if x + dx >= fire_array.shape[0]:
                return fire_array[x, y]
            if y == 0 and dy < 0:
                return fire_array[x, y]
            if y + dy >= fire_array.shape[1]:
                return fire_array[x, y]
            if fire_array[x + dx, y + dy] < np.inf:
                return fire_array[x + dx, y + dy]
            else:
                return fire_array[x, y]

        prop_delta_x = np.zeros_like(fire_array, np.int)
        prop_delta_y = np.zeros_like(fire_array, np.int)
        prop_dir = np.zeros_like(fire_array)

        for x in range(0, fire_array.shape[0]):
            for y in range(0, fire_array.shape[1]):
                if fire_array[x, y] < np.inf:
                    ign = lambda dx, dy: default_ignition(x, y, dx, dy)
                    prop_delta_x[x, y] = ign(1, -1) + 2 * ign(1, 0) + ign(1, 1) - ign(-1,
                                                                                -1) - 2 * ign(
                        -1, 0) - ign(-1, 1)
                    prop_delta_y[x, y] = ign(1, 1) + 2 * ign(0, 1) + ign(-1, 1) - ign(1,
                                                                                -1) - 2 * ign(
                        0, -1) - ign(-1, -1)
                    prop_dir = np.arctan2(prop_delta_y[x, y], prop_delta_x[x, y])

        return prop_delta_x, prop_delta_y, prop_dir

    @staticmethod
    def _compute_traversal_end(fire_array: np.ndarray):
        """Compute the ignition end time for each cell.

        The ignition end time is the latest ignition time among all neighbor cells."""
        end = np.zeros_like(fire_array) * np.inf

        for x in range(0, fire_array.shape[0]):
            for y in range(0, fire_array.shape[1]):
                if fire_array[x, y] < np.inf:
                    max_neighbor = 0.0
                    for dx in [-1, 0, 1]:
                        for dy in [-1, 0, 1]:
                            if (dx == 0 and dy == 0) or (x == 0 and dx < 0) or (
                                    x + dx >= fire_array.shape[0]) or (y == 0 and dy < 0) or (
                                    y + dy >= fire_array.shape[1]):
                                continue
                            if fire_array[x + dx, y + dy] < np.inf:
                                max_neighbor = max(max_neighbor, fire_array[x + dx, y + dy])
                    if max_neighbor <= fire_array[x, y]:
                        # propagation_border
                        end[x, y] = fire_array[x, y] + 180  # assume 3 minutes
                    else:
                        end[x, y] = max_neighbor
                else:
                    end[x, y] = fire_array[x, y]
        return end

    def find_parent_or_child_of_time(self, start_cell: fire_rs.geodata.geo_data.Cell, time: float):
        """Find an ignited cell at 'time' in the propagation graph starting from 'start_cell'

        This function similar to the c++ planning 'project on firefront' algorithm.
        """

        coord = start_cell
        found = False
        while not found:
            if self.geodata.data[self._ignition_layer][coord] <= time and time < \
                    self.geodata.data[self._ignition_end_layer][coord]:
                found = True
            else:
                uphill_cell = (coord[0] + self.geodata.data[self._prop_x_layer][coord],
                               coord[1] + self.geodata.data[self._prop_y_layer][coord])
                downhill_cell = (coord[0] - self.geodata.data[self._prop_x_layer][coord],
                                 coord[1] - self.geodata.data[self._prop_y_layer][coord])
                if uphill_cell[0] == downhill_cell[0] and uphill_cell[1] == downhill_cell[1]:
                    # Extrema reached
                    coord = uphill_cell
                    found = True
                next_cell = downhill_cell if self.geodata.data[self._ignition_layer][
                                                 coord] > time else uphill_cell

                if (next_cell[0] < 0 or next_cell[0] >
                    self.geodata.data[self._ignition_layer].shape[0] - 1) and (
                        next_cell[1] < 0 or next_cell[1] >
                        self.geodata.data[self._ignition_layer].shape[1] - 1):
                    found = True
                if (next_cell == downhill_cell) and (
                        self.geodata.data[self._ignition_layer][next_cell] >
                        self.geodata.data[self._ignition_layer][coord]) or (
                        next_cell == uphill_cell) and (
                        self.geodata.data[self._ignition_layer][next_cell] <
                        self.geodata.data[self._ignition_layer][coord]):
                    # Extrema reached
                    found = True
                if self.geodata.data[self._ignition_layer][next_cell] < np.inf:
                    coord = next_cell
        return coord

def warp_firemap(gd: fire_rs.geodata.geo_data.GeoData, orig: ty.Sequence[ty.Tuple[int, int]],
                 dest: ty.Sequence[ty.Tuple[int, int]],
                 layer: str = "ignition"):
    array = _warp_image(gd[layer], orig, dest)
    new_gd = fire_rs.firemodel.propagation.empty_firemap(gd, layer=layer)
    new_gd.data["ignition"] = array
    return new_gd

def _warp_image(array: np.ndarray, orig: ty.Sequence[ty.Tuple[int, int]],
                dest: ty.Sequence[ty.Tuple[int, int]]) -> np.ndarray:
    newarray = array.copy()
    newarray = newarray.T  # opencv convention on columns and rows is inverted
    newarray_min = newarray.min()
    newarray_max = newarray.max()
    newarray = (newarray - newarray_min) / (newarray_max - newarray_min)
    orig_p = np.array(orig, np.int32)
    orig_p = orig_p.reshape(1, -1, 2)  # Don't ask why. https://stackoverflow.com/a/47114049
    dest_p = np.array(dest, np.int32)
    dest_p = dest_p.reshape(1, -1, 2)  # Don't ask why. https://stackoverflow.com/a/47114049
    good_matches = [cv2.DMatch(p, p, 0) for p in range(len(orig))]
    tps = cv2.createThinPlateSplineShapeTransformer()
    tps.estimateTransformation(dest_p, orig_p, good_matches)
    warped = None
    warped = tps.warpImage(newarray, warped, cv2.INTER_LINEAR, cv2.BORDER_CONSTANT, np.inf)
    warped = warped * (newarray_max - newarray_min) + newarray_min
    warped = warped.T
    return warped