from __future__ import print_function
import numpy as np
import os
import warnings
from collections import defaultdict
from astropy import constants
from astropy import log
from .. import base_class
from ..utils import ImmutableDict,unitless,lower_keys
from .. import utils

import astropy.units as u
_quantity = u.Quantity

class Fjdu(base_class.RadiativeTransferApproximator):
    def __init__(self, datapath=None, species='co',
                 density=None,
                 collider_densities=None,
                 temperature=None,
                 tbg=2.73,
                 column=None,
                 abundance=None,
                 escapeProbGeom='lvg',
                 **kwargs):

        if os.getenv('RADEX_DATAPATH') and datapath is None:
            datapath = os.getenv('RADEX_DATAPATH')

        self.datapath = os.path.dirname(datapath)
        self.species = species

        self._is_locked = False
        self._locked_parameter = None
        self.set_default_params()
        self.set_params(temperature=temperature, density=density,
                        collider_densities=collider_densities,
                        column=column, geotype=escapeProbGeom, **kwargs)
        self.tbg = tbg
        from pyradex.fjdu import wrapper_my_radex
        myradex_wrapper = wrapper_my_radex.myradex_wrapper
        self._myradex = myradex_wrapper
        self._is_locked = False
        self._locked_parameter = None

        if abundance is not None:
            if density is not None or collider_densities is not None:
                self._locked_parameter = 'density'
            elif column is not None:
                self._locked_parameter = 'column'
            else:
                raise ValueError("At least two of column, density, and "
                                 "abundance must be specified")
            self.abundance = abundance

    def __call__(self, return_table=True, **kwargs):

        niter = self.run_radex(**kwargs)

        if return_table:
            return self.get_table()
        else:
            return niter

    def load_datafile(self, filename=None, verbose=False):
        filename = filename or self.molpath
        self.datapath = (os.path.dirname(filename) or self.datapath)+"/"
        self.fname = os.path.basename(filename)

        nlevels, nitems, ntrans = self._myradex.config_basic(self.datapath,
                                                             self.fname,
                                                             unitless(self.tbg),
                                                             verbose)
        self.set_params(**{'n_levels': nlevels,
                           'n_item': nitems,
                           'n_transitions': ntrans})

    def run_radex(self, **kwargs):

        # drop kwargs kept for compatibility with Radex
        ignore_kwargs = ['reuse_last', 'reload_molfile']
        for ik in ignore_kwargs:
            if ik in kwargs:
                kwargs.pop(ik)

        self.set_params(**kwargs)
        self.load_datafile()
        energies, f_occupations, data_transitions, cooling_rate = \
                self._myradex.run_one_params(**self.params)
        self._energies = u.Quantity(energies, u.K) # excitation temperature
        self._data_dict = cast_into_dic(self._myradex.column_names.tostring().decode(),
                                        data_transitions)
        self._level_population = f_occupations

    _default_params = (('tkin', 0.0),
                       ('dv_CGS', 1e5),
                       ('dens_X_CGS', 0.0),
                       ('Ncol_X_CGS', 0.0),
                       ('H2_density_CGS', 0.0),
                       ('HI_density_CGS', 0.0),
                       ('oH2_density_CGS', 0.0),
                       ('pH2_density_CGS', 0.0),
                       ('HII_density_CGS', 0.0),
                       ('Electron_density_CGS', 0.0),
                       ('n_levels', 0),
                       ('n_item', 0),
                       ('n_transitions', 0),
                       ('geotype', 'lvg'),
                      )

    _keyword_map = {#'temperature': 'tkin',
                    'deltav': 'dv_cgs',
                    #'column': 'ncol_x_cgs',
                   }

    _density_keyword_map = {'h2': 'h2_density_cgs',
                            'h': 'hi_density_cgs',
                            'oh2': 'oh2_density_cgs',
                            'ph2': 'ph2_density_cgs',
                            'hii': 'hii_density_cgs',
                            'e': 'electron_density_cgs',
                           }

    def set_default_params(self):
        self._params = lower_keys(dict(self._default_params))

    def set_params(self, **kwargs):
        default = lower_keys(dict(self._default_params))
        for k in kwargs:
            if kwargs[k] is None:
                continue
            if k == 'deltav':
                # deltav requires unit conversion
                self.deltav = kwargs[k]
            elif k.lower() in self._keyword_map:
                self._params[self._keyword_map[k]] = kwargs[k]
            elif k.lower() in ('density','collider_densities'):
                self.density = kwargs[k]
            elif k.lower() == 'column':
                self.column = kwargs[k]
            elif k.lower() == 'temperature':
                # temperature _cannot_ be set until density is
                if not hasattr(self, '_use_thermal_opr'):
                    try:
                        self.density = kwargs['density']
                    except KeyError:
                        self.density = kwargs['collider_densities']
                self.temperature = kwargs[k]
            elif k == 'tbg':
                self.tbg = kwargs[k]
            elif k == 'species':
                self.species = kwargs[k]
            elif k.lower() not in default:
                raise ValueError("{0} is not a valid key.".format(k))
            else:
                self._params[k] = kwargs[k]

    @property
    def params(self):
        return lower_keys(self._params)

    @params.setter
    def params(self, value):
        if not isinstance(value, dict):
            raise TypeError('Parameters must be a dictionary.')
        self.set_params(**value)

    @property
    def density(self):

        dd = {'H2': u.Quantity(self.params['h2_density_cgs'], self._u_cc),
              'OH2': u.Quantity(self.params['oh2_density_cgs'], self._u_cc),
              'PH2': u.Quantity(self.params['ph2_density_cgs'], self._u_cc),
              'E': u.Quantity(self.params['electron_density_cgs'], self._u_cc),
              'H+': u.Quantity(self.params['hii_density_cgs'], self._u_cc),
              'H': u.Quantity(self.params['hi_density_cgs'], self._u_cc),
              'He': u.Quantity(0, self._u_cc),}
        return ImmutableDict(dd)

    @density.setter
    def density(self, collider_density):

        self._use_thermal_opr = False

        if isinstance(collider_density, (float,int,_quantity,np.ndarray)):
            log.warn("Assuming the density is n(H_2).")
            collider_density = {'H2': collider_density}

        collider_densities = defaultdict(lambda: 0)
        for k in collider_density:
            collider_densities[k.upper()] = unitless(u.Quantity(collider_density[k],
                                                                self._u_cc))
            if k.upper() not in self._all_valid_colliders:
                raise ValueError('Collider %s is not one of the valid colliders: %s' %
                                 (k,self._all_valid_colliders))

        if (('OH2' in collider_densities and collider_densities['OH2'] !=0) or
            ('PH2' in collider_densities and collider_densities['PH2'] !=0)):
            if not 'PH2' in collider_densities or not 'OH2' in collider_densities:
                raise ValueError("If o-H2 density is specified, p-H2 must also be.")
            # dictionary of collider densities
            for k in collider_densities:
                if k.lower() in self._density_keyword_map:
                    key = self._density_keyword_map[k.lower()]
                    self._params[key.lower()] = collider_densities[k]
                elif k.lower() in self._density_keyword_map.values():
                    self._params[k.lower()] = collider_densities[k]
                else:
                    raise KeyError("Collider {0} not recognized.".format(k))
            self._params['dens_x_cgs'] = self.total_density.value
            self._use_thermal_opr = False
        elif 'H2' in collider_densities and 'H2' in self._valid_colliders:
            # H2 is a collider in the file: use it.
            self._params['dens_x_cgs'] = collider_density['H2']
            for k in self._density_keyword_map.values():
                self._params[k] = 0.0
            self._params['h2_density_cgs'] = collider_density['H2']
        elif 'H2' in collider_densities:
            # Only oH2 and pH2 are in the file.  Must assume.
            warnings.warn("Using a default ortho-to-para ratio (which "
                          "will only affect species for which independent "
                          "ortho & para collision rates are given)")
            self._use_thermal_opr = True
            #self.radex.cphys.density[0] = collider_densities['H2']

            T = unitless(self.temperature)
            if T > 0:
                # From Faure, private communication
                opr = min(3.0,9.0*np.exp(-170.6/T))
            else:
                opr = 3.0
            fortho = opr/(1+opr)
            log.debug("Set OPR to {0} and fortho to {1}".format(opr,fortho))
            self._params['oh2_density_cgs'] = collider_density['H2']*(fortho)
            self._params['ph2_density_cgs'] = collider_density['H2']*(1-fortho)
            self._params['dens_x_cgs'] = self.total_density.value

    @property
    def temperature(self):
        return u.Quantity(self.params['tkin'], u.K)

    @temperature.setter
    def temperature(self, tkin):
        if hasattr(tkin,'to'):
            tkin = unitless(u.Quantity(tkin, u.K))
        if tkin <= 0 or tkin > 1e4:
            raise ValueError('Must have kinetic temperature > 0 and < 10^4 K')
        self.set_params(tkin=tkin)

        if self._use_thermal_opr:
            # Reset the density to a thermal value
            lp = self._locked_parameter
            self.density = (unitless(self.density['H2']) or
                            unitless(self.density['OH2']+self.density['PH2']))
            self._locked_parameter = lp

    @property
    def column_per_bin(self):
        return u.Quantity(self.params['ncol_x_cgs'], self._u_sc)

    @column_per_bin.setter
    def column_per_bin(self, col):
        if hasattr(col,'to'):
            col = unitless(u.Quantity(col, self._u_sc))
        if col < 1e5 or col > 1e25:
            raise ValueError("Extremely low or extremely high column.")
        self.set_params(ncol_x_cgs=col)

        col = u.Quantity(col, self._u_sc)
        if not self._is_locked:
            self._is_locked = True
            if self.locked_parameter == 'density':
                ab = (col/(self.total_density * self.length))
                if hasattr(ab, 'decompose'):
                    self.abundance = ab.decompose().value
                else:
                    self.abundance = ab / (self._u_cc*u.pc).to(self._u_sc)
            elif self.locked_parameter == 'abundance':
                self.density = col / self.length / self.abundance
            self._lock_param('column')
            self._is_locked = False

    @property
    def abundance(self):
        return self._abundance

    @abundance.setter
    def abundance(self, abund):
        self._abundance = abund
        if not self._is_locked:
            self._is_locked = True
            if self.locked_parameter == 'column':
                dens = self.column_per_bin / self.length / abund
                self.density = dens
            elif self.locked_parameter == 'density':
                col = self.total_density*self.length*abund
                self.column_per_bin = u.Quantity(col, u.cm**-2)
            self._lock_param('abundance')
            self._is_locked=False

    @property
    def tbg(self):
        return u.Quantity(self._tbg, u.K)

    @tbg.setter
    def tbg(self, tbg):
        if hasattr(tbg, 'value'):
            self._tbg = unitless(u.Quantity(tbg, u.K))
        else:
            self._tbg = tbg

    @property
    def deltav(self):
        return u.Quantity(self.params['dv_cgs']*self._kms_to_cms, self._u_kms)

    _kms_to_cms = 1e-5
    _u_cms = u.cm/u.s

    @deltav.setter
    def deltav(self, dv):
        if hasattr(dv, 'unit'):
            self._params['dv_cgs'] = unitless(dv.to(self._u_cms))
        else:
            self._params['dv_cgs'] = unitless(u.Quantity(dv/self._kms_to_cms,
                                                         self._u_cms))

    @property
    def molpath(self):
        if hasattr(self,'_molpath'):
            return self._molpath

    @molpath.setter
    def molpath(self, molfile):
        if "~" in molfile:
            molfile = os.path.expanduser(molfile)
        utils.verify_collisionratefile(molfile)
        self._molpath = molfile

    @property
    def datapath(self):
        return self._datapath

    @datapath.setter
    def datapath(self, datapath):
        self._datapath = datapath

    @property
    def escapeprobProbGeom(self):
        return self._params['geotype']

    @escapeprobProbGeom.setter
    def escapeprobProbGeom(self, value):
        if value in ('lvg','spherical','slab'):
            self._params['geotype'] = value
        else:
            raise ValueError("Geometry must be spherical, slab, or lvg")

    _um_to_ghz = u.um.to(u.GHz, equivalencies=u.spectral())

    @property
    def frequency(self):
        return u.Quantity(self._um_to_ghz/self._data_dict['lam'], unit=u.GHz)

    @property
    def level_population(self):
        return self._level_population

    @property
    def tex(self):
        return u.Quantity(self._data_dict['Tex'], u.K)

    Tex = tex

    @property
    def tau(self):
        return self._data_dict['tau']

    @property
    def upperstateenergy(self):
        return u.Quantity(self._data_dict['Eup'], u.K)

    @property
    def upperlevelnumber(self):
        return self._data_dict['iup']

    @property
    def lowerlevelnumber(self):
        return self._data_dict['ilow']

    @property
    def upperlevelpop(self):
        return self._data_dict['fup']

    @property
    def lowerlevelpop(self):
        return self._data_dict['flow']

    @property
    def source_line_brightness_temperature(self):
        return u.Quantity(self._data_dict['Tr'], u.K)

    #@property
    #def source_line_surfbrightness(self):
    #    return u.Quantity(self._data_dict['flux'], self._u_brightness)

    @property
    def source_brightness(self):
        return u.Quantity(self._data_dict['flux_dens'], self._u_brightness)

    @property
    def background_brightness(self):
        return u.Quantity(self._data_dict['Jback'], self._u_brightness)
    #    return self.tbg.to(self._u_brightness)

    @property
    def beta(self):
        return self._data_dict['beta']

    @property
    def statistical_weight(self):
        return self._data_dict['gup']

def cast_into_dic(col_names, arr):
    '''col_names is column_info, and arr is data_transitions'''
    names = col_names.split()
    return {names[i]: arr[i,:] for i in range(len(names))}

import numpy,import astropy,from astropy


import os, sys, glob, copy
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Union, Tuple, TypeVar
from importlib import import_module

import pandas as pd
import numpy as np
import numba as nb
from skimage import transform

from joblib import Parallel, delayed
from multiprocessing import cpu_count

import astropy
from astropy.constants import sigma_T, m_p, c
from astropy.io import fits
from astropy import units as un
import lenstools
from lenstools import ConvergenceMap
#import pymaster as nmt

from astrild.rays.utils import Filters
from astrild.rays.skys.sky_utils import SkyUtils, SkyNumbaUtils
from astrild.rays.skyio import SkyIO
from astrild.io import IO

dir_src = Path(__file__).parent.parent.absolute()
default_config_file_ray = dir_src / "configs/ray_snapshot_info.h5"

sigma_T = sigma_T.to(un.Mpc**2).value #[Mpc^2]
m_p = m_p.to(un.M_sun).value #[M_sun]
c_light = c.to("km/s").value
T_cmb = 2.7251 #[K]
Gcm2 = 4.785E-20 # G/c^2 (Mpc/M_sun)

# store available nr. of cpus for parallel computation
ncpus_available = cpu_count()

class SkyArrayWarning(BaseException):
    pass

class SkyArray:
    """
    The sky-map contains either:
        - results of ray-tracing simulations run with RayRamses
        - results of analytical dT and deflection angle results for NFW halos

    This class analyzes the 2D map that contains the summes pertrurbations
    of each ray. It can prepare the data for the search of voids and peaks.

    Attributes:
        skymap:
        opening_angle: [deg]
        quantity:
        dirs:
        map_file:

    Methods:
        from_file:
        from_dataframe:
        from_array:
        from_halo_series:
        from_halo_dataframe:
        pdf:
        wl_peak_counts:
        add_galaxy_shape_noise:
        create_galaxy_shape_noise:
        create_cmb:
        create_mask:
        convolution:
        zoom:
        division:
        merge:
        convert_convergence_to_deflection:
    """

    def __init__(
        self,
        skymap: np.ndarray,
        opening_angle: float,
        quantity: str,
        dirs: Dict[str, str],
        map_file: Optional[str] = None,
    ):
        self.data = {"orig": skymap}
        self._npix = skymap.shape[0]
        self._opening_angle = opening_angle
        self.quantity = quantity
        self.dirs = dirs
        self.map_file = map_file

    @classmethod
    def from_file(
        cls,
        map_file: str,
        opening_angle: float,
        quantity: str,
        dir_in: str,
        npix: Optional[int] = None,
        convert_unit: bool = True,
    ) -> "SkyArray":
        """
        Initialize class by reading the skymap data from pandas hdf5 file
        or numpy array.
        The file can be pointed at via map_filename or file_dsc.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            opening_angle: [deg]
        """
        assert map_file, "There is no file being pointed at"

        file_extension = map_file.split(".")[-1]
        if file_extension == "h5":
            map_df = pd.read_hdf(map_file, key="df")
            return cls.from_dataframe(
                map_df,
                opening_angle,
                quantity,
                dir_in,
                map_file,
                npix,
                convert_unit,
            )
        elif file_extension in ["npy", "fits"]:
            if file_extension == "npy":
                map_array = np.load(map_file)
            elif file_extension == "fits":
                map_array = ConvergenceMap.load(map_file, format="fits").data
            return cls.from_array(
                map_array, opening_angle, quantity, dir_in, map_file
            )

    @classmethod
    def from_dataframe(
        cls,
        map_df: pd.DataFrame,
        opening_angle: float,
        quantity: str,
        dir_in: str,
        map_file: str,
        npix: Optional[int] = None,
        convert_unit: bool = True,
    ) -> "SkyArray":
        """
        Initialize class by reading the skymap data from pandas DataFrame.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            opening_angle: [deg]
        """
        if convert_unit:
            map_df = SkyUtils.convert_code_to_phy_units(quantity, map_df)
        map_array = SkyIO.transform_RayRamsesOutput_to_NumpyNdarray(
            map_df[quantity].values
        )
        return cls.from_array(
            map_array, opening_angle, quantity, dir_in, map_file
        )

    @classmethod
    def from_array(
        cls,
        map_array: np.array,
        opening_angle: float,
        quantity: str,
        dir_in: str,
        map_file: Optional[str] = None,
    ) -> "SkyArray":
        """
        Initialize class by reading the skymap data from np.ndarray.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            opening_angle: [deg]
        """
        assert map_array.shape[0] == map_array.shape[1]
        dirs = {"sim": dir_in}
        return cls(map_array, opening_angle, quantity, dirs, map_file)

    @classmethod
    def from_halo_series(
        cls,
        halo: pd.Series,
        npix: int = 100,
        extent: float = 1,
        direction: List[int] = [0, 1],
        suppress: bool = False,
        suppression_R: float = 1,
        to: str = "dT",
    ) -> "SkyArray":
        """
        Calculate the deflection angle of Rees-Sciama / Birkinshaw-Gull /
        moving cluster of galaxies effect of a halo with NFW profile using method
        in described Sec. 3.2 in Baxter et al 2015 (1412.7521).

        Note:
            In this application it can be assumed that s_{SL}/s_{S}=1. Furthermore,
            we can deglect vec{theta}/norm{theta} as it will be multiplied by the
            same in the integral of Eq. 9 in Yasini et a. 2018 (1812.04241).

        Args:
            theta_200c: radius, [deg]
            M_200c: mass, [Msun]
            c_200c: concentration, [-]
            extent: The size of the map from which the trans-vel is calculated
                in units of R200 of the associated halo.
            suppress:
            suppression_R:
            angu_diam_dist: angular diameter distance, [Mpc]
            direction: 0=(along x-axis), 1=(along y-axis), if 0 and 1 are given
                the sum of both maps is returned.
            to: Indicates whether 2nd order temperature perturbation of Ree-Sciama
                effect or deflection angle map will be created, [dT, alpha]
        """
        halo_dict = {
            "theta_200c": halo.r200_deg,
            "M_200c": halo.m200,
            "c_200c": halo.c_NFW,
            "angu_diam_dist": halo.Dc,
        }

        # set quantity label of SkyArray and method variable
        if to == "dT":
            quantity = "rs"
            get_analytical_map = SkyUtils.NFW_temperature_perturbation_map
            halo_dict["vel"] = [halo.theta1_tv, halo.theta2_tv]
        elif to == "alpha":
            quantity = "alpha"
            get_analytical_map = SkyUtils.NFW_deflection_angle_map
        else:
            SkyArrayWarning("The routine for this quantity is not implemented")

        # place directional indicator in SkyArray quantity label
        if 1 in direction and 0 in direction:
            pass
        elif 0 in direction:
            quantity += "_x"
        else:
            quantity += "_y"

        map_array = get_analytical_map(
            **halo_dict,
            npix=npix,
            extent=extent,
            direction=direction,
            suppress=suppress,
            suppression_R=suppression_R,
        )
        opening_angle = 2 * halo.r200_deg * extent
        return cls(map_array, opening_angle, quantity, dirs=None, map_file=None)

    @classmethod
    def from_halo_dataframe(
        cls,
        halo_cat: Union[pd.DataFrame, pd.Series],
        npix: int = 8192,
        extent: float = 1,
        direction: List[int] = [0, 1],
        suppress: bool = False,
        suppression_R: float = 1,
        opening_angle: float = 20.,
        ncpus: int = 1,
        to: str = "dT",
    ) -> "SkyArray":
        """
        Identical to SkyArray.from_halo_series() but for a pd.DataFrame.
        """
        halo_dict = halo_cat[[
            "r200_deg",
            "r200_pix",
            "m200",
            "c_NFW",
            "Dc",
            "theta1_pix",
            "theta2_pix",
        ]].to_dict(orient='list')
        halo_idx = np.arange(len(halo_dict["m200"]))

        # set quantity label of SkyArray and method variable
        if to == "dT":
            quantity = "rs"
            halo_dict["theta1_tv"] = halo_cat["theta1_tv"].values
            halo_dict["theta2_tv"] = halo_cat["theta2_tv"].values
        elif to == "alpha":
            quantity = "alpha"
        else:
            SkyArrayWarning("The routine for this quantity is not implemented")

        # place directional indicator in SkyArray quantity label
        if 1 in direction and 0 in direction:
            pass
        elif 0 in direction:
            quantity += "_x"
        else:
            quantity += "_y"

        if ncpus == 1:
            map_array = SkyUtils.analytic_Halo_signal_to_SkyArray(
                halo_idx,
                halo_dict,
                extent,
                direction,
                suppress,
                suppression_R,
                npix,
                to,
            )
        else:
            halo_idx_batches = np.array_split(halo_idx, ncpus)
            map_sub_arrays = Parallel(n_jobs=ncpus)(
                delayed(SkyUtils.analytic_Halo_signal_to_SkyArray)(
                    halo_idx_batch,
                    halo_dict,
                    extent,
                    direction,
                    suppress,
                    suppression_R,
                    npix,
                    to,
                ) for halo_idx_batch in halo_idx_batches
            )
            map_array = sum(map_sub_arrays)
        map_array = np.nan_to_num(
            map_array, copy=False, nan=0.0, posinf=0.0, neginf=0.0,
        )

        return cls(map_array, opening_angle, quantity, dirs=None, map_file=None)

    @classmethod
    def from_halo_catalogue_to_temperature_perturbation_map(
        cls,
        halo_cat: pd.DataFrame,
        extent: float = 1,
        direction: List[int] = [0, 1],
        suppress: bool = False,
        suppression_R: float = 1,
        npix: int = 8192,
        opening_angle: float = 20.,
        ncpus: int = 1,
    ) -> "SkyArray":
        """
        The Rees-Sciama / Birkinshaw-Gull / moving cluster of galaxies effect.

        Args:
            vel: transverse to the line-of-sight velocity, [km/sec]

        Returns:
            Temperature perturbation map, \Delta T / T_CMB
        """
        halo_dict = halo_cat[[
            "r200_deg",
            "r200_pix",
            "m200",
            "c_NFW",
            "Dc",
            "theta1_pix",
            "theta2_pix",
            "theta1_tv",
            "theta2_tv",
        ]].to_dict(orient='list')
        halo_idx = range(len(halo_dict["m200"]))

        if ncpus == 1:
            map_array = SkyUtils.analytic_Halo_signal_to_SkyArray(
                halo_idx, halo_dict, extent, direction, suppress, suppression_R, npix
            )
        else:
            halo_idx_batches = np.array_split(halo_idx, ncpus)
            map_sub_arrays = Parallel(n_jobs=ncpus)(
                delayed(SkyUtils.analytic_Halo_signal_to_SkyArray)(
                    halo_idx_batch,
                    halo_dict,
                    extent,
                    direction,
                    suppress,
                    suppression_R,
                    npix,
                ) for halo_idx_batch in halo_idx_batches
            )
            map_array = sum(map_sub_arrays)
        map_array[np.isinf(map_array)] = 0.

        if 1 in direction and 0 in direction:
            quantity = "isw_rs"
        elif 0 in direction:
            quantity = "isw_rs_x"
        else:
            quantity = "isw_rs_y"
        return cls(map_array, opening_angle, quantity, dirs=None, map_file=None)

    @property
    def ncpus(self):
        return self._ncpus

    @ncpus.setter
    def ncpus(self, val: int):
        if (ncpus == 0) or (ncpus < -1):
            raise ValueError(
                f"ncpus={ncpus} is not valid. Please enter a value " +\
                ">0 for ncpus or -1 to use all available cores."
            )
        elif ncpus == -1: self._ncpus = ncpus_available
        else: self._ncpus = val

    @property
    def npix(self) -> int:
        return self._npix

    @property
    def opening_angle(self) -> float:
        return self._opening_angle

    def pdf(self, nbins: int, of: str = "orig") -> dict:
        _pdf = {}
        _pdf["values"], _pdf["bins"] = np.histogram(
            self.data[of], bins=nbins, density=True
        )
        return _pdf

    def wl_peak_counts(
        self,
        nbins: int,
        field_conversion: str,
        of: str = "orig",
        limits: Optional[tuple] = None,
    ) -> pd.DataFrame:
        """
        Signal peak counts. This is used commonly used in weak-lensing,
        but it doesn't need to stop there...
        """
        if field_conversion == "normalize":
            _map = self.data[of] - np.mean(self.skymap.data[of])
        else:
            _map = self.data[of]

        if limits is None:
            lower_bound = np.percentile(
                self.data[of], 5
            )  # np.min(self.data[of])
            upper_bound = np.percentile(
                self.data[of], 95
            )  # np.max(self.data[of])
        else:
            lower_bound = min(limits)
            upper_bound = max(limits)

        map_bins = np.arange(
            lower_bound, upper_bound, (upper_bound - lower_bound) / nbins
        )
        _map = ConvergenceMap(data=_map, angle=self._opening_angle * un.deg)
        _kappa, _pos = _map.locatePeaks(map_bins)

        _hist, _kappa = np.histogram(_kappa, bins=nbins, density=False)
        _kappa = (_kappa[1:] + _kappa[:-1]) / 2
        peak_counts_dic = {"kappa": _kappa, "counts": _hist}
        peak_counts_df = pd.DataFrame(data=peak_counts_dic)
        return peak_counts_df

    def resize(
        self,
        npix,
        of: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Union[np.ndarray, None]:
        """
        Lower the nr. of pixels of image. Useful for tests.

        Args:
            npix: the new pixel nr. per edge of the image
            of: skymap image identifier.
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy
        """
        img = self._manage_img_data(img, orig_data)
        img = transform.resize(img, (npix, npix), anti_aliasing=True)
        if rtn: return img
        else: self.data[of] = img

    def crop(
        self,
        xlimit: Union[Tuple, List, np.array],
        ylimit: Union[Tuple, List, np.array],
        of: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Union[np.ndarray, None]:
        """
        Zoom into sky_array map.

        Args:
            x,ylimit: Boundaries of zoom. If given in ints units are pixels,
                if floats percentages are used.
            of: skymap image identifier.
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy
        """
        if of:
            assert of in list(self.data.keys()), "Map does not exist."
            img = self.data[of]

        img = self._manage_img_data(img, orig_data)
        xlimit = np.asarray(xlimit)
        ylimit = np.asarray(ylimit)
        assert np.diff(xlimit) == np.diff(ylimit), SkyArrayWarning("The whole class is currently designed for square images.")
        if isinstance(xlimit[0], float):
            _npix = img.shape[0]
            xlimit = (_npix * xlimit / 100).astype(int)
            ylimit = (_npix * ylimit / 100).astype(int)
        zoom = img[xlimit[0] : xlimit[1], ylimit[0] : ylimit[1]]
        if rtn:
            return zoom
        else:
            print(f"Image crop to x={xlimit} and y={ylimit}.")
            self.data[of] = zoom
            self._opening_angle = (
                self._opening_angle * abs(np.diff(xlimit)) / self._npix
            )
            self._npix = zoom.shape[0]

    def division(
        self,
        ntiles: int,
        of: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Union[List[np.ndarray], None]:
        """
        Divide image into tiles.
        Should use sklearn.feature_extraction.image.extract_patches_2d

        Args:
            ntiles: Nr. of tiles per edge (as to be in 2^n).
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy
        """
        img = self._manage_img_data(img, orig_data)
        npix = img.shape[0]
        edges = list(np.arange(0, npix, npix / ntiles)) + [npix]
        edges = np.array(
            [edges[idx : idx + 2] for idx in range(len(edges) - 1)]
        ).astype(int)
        tiles = []
        for xlim in edges:
            for ylim in edges:
                tiles.append(self.crop(xlim, ylim, img=img, rtn=True))
        print(
            f"The image is divided into {len(tiles)} tiles, " +\
            f"each with {tiles[0].shape[0]}^2 pixels."
        )
        tiles = np.asarray(tiles)
        if rtn:
            return tiles
        else:
            self.tiles = tiles
            self._tile_npix = tiles[0].shape[0]
            self._tile_opening_angle = self._opening_angle * self._tile_npix / self._npix

    def merge(
        self, tiles: np.ndarray, rtn: bool = False
    ) -> Union[np.ndarray, None]:
        """
        Merge tiles created with self.division.

        Args:
            tiles: 3D
        """
        ntiles = len(tiles)
        nrows = int(np.sqrt(ntiles))
        _parts = np.arange(0, ntiles + nrows, nrows)
        row_tile_idx = [(_parts[ii], _parts[ii + 1]) for ii in range(nrows)]
        row_tiles = []
        for idx in range(nrows):
            start = row_tile_idx[idx][0]
            end = row_tile_idx[idx][1]
            row_tiles.append(np.hstack((tiles[start:end])))
        row_tiles = np.asarray(row_tiles)
        img = np.vstack((row_tiles))
        if rtn: return img

    def substract_mean(
        self,
        of: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Union[np.ndarray, None]:
        """
        Centre values of map on it's mean value.
        """
        if of:
            assert of in list(self.data.keys()), "Map does not exist."
            img = self.data[of]
        img = self._manage_img_data(img, orig_data)
        img -= np.mean(img)
        if rtn: return img
        else: self.data[of] = img

    def filter(
        self,
        filter_dsc: dict,
        on: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Union[np.ndarray, None]:
        """
        Apply kernel (filter_dsc) over skymap image (on).

        Args:
            filter_dsc: Kernel description.
                Note that theta_i should be given in astropy.units!
            on: skymap image identifier.
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy
        """
        if on:
            assert on in list(self.data.keys()), "Map does not exist."
            img = self.data[on]
            map_name = [on]
        else:
            map_name = [""]

        # load rays/utils/filters.py package for dynamic function call
        module = import_module("astrild.rays.utils")
        img = self._manage_img_data(img, orig_data)

        for filter_name, args in filter_dsc.items():
            if rtn is False:
                abbrev = args["abbrev"]
                del args["abbrev"]
                map_name.append(abbrev)

            clas = getattr(module, "Filters")
            fct = getattr(clas, filter_name)
            img = fct(img, self._opening_angle * un.deg, **args)
        if rtn: return img
        else: self.data[("_").join(map_name)] = img

    def create_galaxy_shape_noise(
        self, std: float, ngal: float, rnd_seed: Optional[int] = None
    ) -> None:
        """
        Galaxy Shape Noise (GSN), e.g.: arxiv:1907.06657

        Args:
            std: dispersion of source galaxy intrinsic ellipticity, 0.4 for LSST
            ngal: Nr. density of galaxies, 40 for LSST; [arcmin^2]
            rnd_seed: Fix random seed, for reproducability.
        Returns:
            gsn_map:
                self.npix x self.npix np.array containing the GSN
        """
        theta_pix = 60 * self._opening_angle / self._npix
        std_pix = 0.007  # np.sqrt(std ** 2 / (2*theta_pix*ngal))
        if rnd_seed is None:
            self.data["gsn"] = np.random.normal(
                loc=0, scale=std_pix, size=[self._npix, self._npix]
            )
        else:
            rg = np.random.Generator(np.random.PCG64(rnd_seed))
            self.data["gsn"] = rg.normal(
                loc=0, scale=std_pix, size=[self._npix, self._npix]
            )
        print(f"The GSN map sigma is {np.std(self.data['gsn'])}", std_pix)

    def add_galaxy_shape_noise(self, on: str = "orig") -> np.ndarray:
        """
        Add GSN on top of skymap.

        Args:
            std: dispersion of source galaxy intrinsic ellipticity, 0.4 for LSST
            ngal: Nr. density of galaxies, 40 for LSST; [arcmin^2]
            rnd_seed: Fix random seed, for reproducability.
        """
        if "kappa" in self.quantity:
            self.data["orig_gsn"] = self.data["orig"] + self.data["gsn"]
            return self.data["orig_gsn"]
        else:
            raise SkyArrayWarning(f"GSN should not be added to {self.quantity}")

    def create_cmb(
        self,
        filepath_cl: str,
        lmax: int = 3e3,
        rnd_seed: Optional[int] = None,
        rtn: bool = False,
    ) -> np.ndarray:
        """
        Cosmig Microwave Background (CMB) on partial-sky map,
        for which the flat-sky approximation holds (ell > 10).

        Args:
            filepath_cl: angular power spectrum of CMB
            theta: Edge length of the square field-of-view [deg]
            nside: Nr. of pixels per edge of the output full-sky map
            rnd_seed: Fix random seed, for reproducability.

        Returns:
            cmb_map:
        """
        if rnd_seed:
            np.random.seed(rnd_seed)
        Nx = Ny = self._npix
        Lx = Ly = self._opening_angle * np.pi / 180.0

        cl_tt_cmb = np.load(filepath_cl)[1]
        #cmb = nmt.synfast_flat(
        #    Nx, Ny, Lx, Ly, cls=np.array([cl_tt_cmb]), spin_arr=np.array([0])
        #)[0]
        if rtn: return cmb
        else: self.data["cmb"] = cmb

    def add_cmb(
        self,
        filepath_cl: Optional[str] = None,
        filepath_cmb: Optional[str] = None,
        on: str = "orig",
        lmax: Optional[float] = None,
        rnd_seed: Optional[int] = None,
        rtn: bool = False,
        overwrite: bool = True,
    ) -> np.ndarray:
        """
        Args:
            filepath_cl:
            on:
            lmax:
            rnd_seed:
            rtn:
            overwrite:

        Returns:
        """
        if "isw" in self.quantity:
            if "cmb" not in self.data.keys():
                try: self.create_cmb(filepath_cl, lmax, rnd_seed)
                except:
                    self.data["cmb"] = np.load(filepath_cmb)
            _map = self.data[on] + self.data["cmb"]
            if rtn:
                return _map
            else:
                if overwrite:
                    self.data[on] = _map
                else:
                    self.data[f"{on}_cmb"] = _map
        else:
            raise SkyArrayWarning(f"CMB should not be added to {self.quantity}")

    def convert_convergence_to_deflection(
        self,
        on: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        npix: Optional[int] = None,
        opening_angle: Optional[float] = None,
        rtn: bool = True,
        orig_data: str = None,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Args:
            on: String to indicate which map in self.data should be used.
            img: 2D convergence map.
            rtn: Bool to indicate whether to attach result to class object or return.
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy

        Returns:
            alpha_1,2: 2D deflection angle map [rad]
        """
        #TODO handle tiles/multiple images
        assert self.quantity in [
            "kappa_1",
            "kappa_2",
        ], "Deflection angle can only be calculated from the kappa map"
        img = self._manage_img_data(img, orig_data)

        if npix is None: npix = self._npix
        if opening_angle is None: opening_angle = self._opening_angle

        alpha_1, alpha_2 = SkyUtils.convert_convergence_to_deflection_ctypes(
            img, npix, opening_angle * un.deg
        )
        if rtn:
            return alpha_2, alpha_1
        else:
            self.data["defltx"] = alpha_2
            self.data["deflty"] = alpha_1

    def convert_deflection_to_shear(
        self,
        on: Optional[str] = None,
        img: Optional[np.ndarray] = None,
        rtn: bool = False,
        orig_data: str = None,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Args:
            on: String to indicate which map in self.data should be used.
            img: 2D convergence map.
            rtn: Bool to indicate whether to attach result to class object or return.
            orig_data: What to do with data of the image to be processed:
                e.g. no, shallow, or deep copy

        Returns:
            gamma_1,2: 2D shear map [-]
        """
        assert self.quantity in [
            "alpha"
        ], "Shear can only be calculated from the deflection angle map"
        img = self._manage_img_data(img, orig_data)
        gamma_1, gamma_2 = SkyUtils.convert_deflection_to_shear(
            img, self._npix, self._opening_angle * un.deg
        )
        if rtn:
            return gamma_2, gamma_1
        else:
            self.data["gammax"] = gamma_2
            self.data["gammay"] = gamma_1

    @staticmethod
    def _manage_img_data(
        img: np.ndarray,
        orig_data: str = None,
    ) -> np.ndarray:
        """
        Handle the memory location of the image to be processed.

        Args:
            img: memory pointer of the image data.
            orig_data: action key word.

        Returns:
            cimg: pointer to (new) memory location.
        """
        if orig_data == "shallow": cimg = copy.copy(img)
        elif orig_data == "deep": cimg = copy.deepcopy(img)
        else: cimg = img
        return cimg

import numpy,import numba,import astropy,from astropy


from __future__ import absolute_import, division, print_function

from ..compatibility import Sequence
import inspect

import numpy as np

try:
    import scipy
    import scipy.fftpack
except ImportError:
    scipy = None

from .core import concatenate as _concatenate
from .creation import arange as _arange
from ..utils import derived_from

chunk_error = ("Dask array only supports taking an FFT along an axis that \n"
               "has a single chunk. An FFT operation was tried on axis %s \n"
               "which has chunks %s. To change the array's chunks use "
               "dask.Array.rechunk.")

fft_preamble = """
    Wrapping of %s

    The axis along which the FFT is applied must have a one chunk. To change
    the array's chunking use dask.Array.rechunk.

    The %s docstring follows below:

    """

def _fft_out_chunks(a, s, axes):
    """ For computing the output chunks of [i]fft*"""
    if s is None:
        return a.chunks
    chunks = list(a.chunks)
    for i, axis in enumerate(axes):
        chunks[axis] = (s[i],)
    return chunks

def _rfft_out_chunks(a, s, axes):
    """ For computing the output chunks of rfft*"""
    if s is None:
        s = [a.chunks[axis][0] for axis in axes]
    s = list(s)
    s[-1] = s[-1] // 2 + 1
    chunks = list(a.chunks)
    for i, axis in enumerate(axes):
        chunks[axis] = (s[i],)
    return chunks

def _irfft_out_chunks(a, s, axes):
    """ For computing the output chunks of irfft*"""
    if s is None:
        s = [a.chunks[axis][0] for axis in axes]
        s[-1] = 2 * (s[-1] - 1)
    chunks = list(a.chunks)
    for i, axis in enumerate(axes):
        chunks[axis] = (s[i],)
    return chunks

def _hfft_out_chunks(a, s, axes):
    assert len(axes) == 1

    axis = axes[0]

    if s is None:
        s = [2 * (a.chunks[axis][0] - 1)]

    n = s[0]

    chunks = list(a.chunks)
    chunks[axis] = (n,)
    return chunks

def _ihfft_out_chunks(a, s, axes):
    assert len(axes) == 1

    axis = axes[0]

    if s is None:
        s = [a.chunks[axis][0]]
    else:
        assert len(s) == 1

    n = s[0]

    chunks = list(a.chunks)
    if n % 2 == 0:
        m = (n // 2) + 1
    else:
        m = (n + 1) // 2
    chunks[axis] = (m,)
    return chunks

_out_chunk_fns = {'fft': _fft_out_chunks,
                  'ifft': _fft_out_chunks,
                  'rfft': _rfft_out_chunks,
                  'irfft': _irfft_out_chunks,
                  'hfft': _hfft_out_chunks,
                  'ihfft': _ihfft_out_chunks}

def fft_wrap(fft_func, kind=None, dtype=None):
    """ Wrap 1D, 2D, and ND real and complex FFT functions

    Takes a function that behaves like ``numpy.fft`` functions and
    a specified kind to match it to that are named after the functions
    in the ``numpy.fft`` API.

    Supported kinds include:

        * fft
        * fft2
        * fftn
        * ifft
        * ifft2
        * ifftn
        * rfft
        * rfft2
        * rfftn
        * irfft
        * irfft2
        * irfftn
        * hfft
        * ihfft

    Examples
    --------
    >>> parallel_fft = fft_wrap(np.fft.fft)
    >>> parallel_ifft = fft_wrap(np.fft.ifft)
    """
    if scipy is not None:
        if fft_func is scipy.fftpack.rfft:
            raise ValueError("SciPy's `rfft` doesn't match the NumPy API.")
        elif fft_func is scipy.fftpack.irfft:
            raise ValueError("SciPy's `irfft` doesn't match the NumPy API.")

    if kind is None:
        kind = fft_func.__name__
    try:
        out_chunk_fn = _out_chunk_fns[kind.rstrip("2n")]
    except KeyError:
        raise ValueError("Given unknown `kind` %s." % kind)

    def func(a, s=None, axes=None):
        if axes is None:
            if kind.endswith('2'):
                axes = (-2, -1)
            elif kind.endswith('n'):
                if s is None:
                    axes = tuple(range(a.ndim))
                else:
                    axes = tuple(range(len(s)))
            else:
                axes = (-1,)
        else:
            if len(set(axes)) < len(axes):
                raise ValueError("Duplicate axes not allowed.")

        _dtype = dtype
        if _dtype is None:
            sample = np.ones(a.ndim * (8,), dtype=a.dtype)
            try:
                _dtype = fft_func(sample, axes=axes).dtype
            except TypeError:
                _dtype = fft_func(sample).dtype

        for each_axis in axes:
            if len(a.chunks[each_axis]) != 1:
                raise ValueError(chunk_error % (each_axis, a.chunks[each_axis]))

        chunks = out_chunk_fn(a, s, axes)

        args = (s, axes)
        if kind.endswith('fft'):
            axis = None if axes is None else axes[0]
            n = None if s is None else s[0]
            args = (n, axis)

        return a.map_blocks(fft_func, *args, dtype=_dtype,
                            chunks=chunks)

    if kind.endswith('fft'):
        _func = func

        def func(a, n=None, axis=None):
            s = None
            if n is not None:
                s = (n,)

            axes = None
            if axis is not None:
                axes = (axis,)

            return _func(a, s, axes)

    func_mod = inspect.getmodule(fft_func)
    func_name = fft_func.__name__
    func_fullname = func_mod.__name__ + "." + func_name
    if fft_func.__doc__ is not None:
        func.__doc__ = (fft_preamble % (2 * (func_fullname,)))
        func.__doc__ += fft_func.__doc__
    func.__name__ = func_name
    return func

fft = fft_wrap(np.fft.fft)
fft2 = fft_wrap(np.fft.fft2)
fftn = fft_wrap(np.fft.fftn)
ifft = fft_wrap(np.fft.ifft)
ifft2 = fft_wrap(np.fft.ifft2)
ifftn = fft_wrap(np.fft.ifftn)
rfft = fft_wrap(np.fft.rfft)
rfft2 = fft_wrap(np.fft.rfft2)
rfftn = fft_wrap(np.fft.rfftn)
irfft = fft_wrap(np.fft.irfft)
irfft2 = fft_wrap(np.fft.irfft2)
irfftn = fft_wrap(np.fft.irfftn)
hfft = fft_wrap(np.fft.hfft)
ihfft = fft_wrap(np.fft.ihfft)

def _fftfreq_block(i, n, d):
    r = i.copy()
    r[i >= (n + 1) // 2] -= n
    r /= n * d
    return r

@derived_from(np.fft)
def fftfreq(n, d=1.0, chunks='auto'):
    n = int(n)
    d = float(d)

    r = _arange(n, dtype=float, chunks=chunks)

    return r.map_blocks(_fftfreq_block, dtype=float, n=n, d=d)

@derived_from(np.fft)
def rfftfreq(n, d=1.0, chunks='auto'):
    n = int(n)
    d = float(d)

    r = _arange(n // 2 + 1, dtype=float, chunks=chunks)
    r /= n * d

    return r

def _fftshift_helper(x, axes=None, inverse=False):
    if axes is None:
        axes = list(range(x.ndim))
    elif not isinstance(axes, Sequence):
        axes = (axes,)

    y = x
    for i in axes:
        n = y.shape[i]
        n_2 = (n + int(inverse is False)) // 2

        l = y.ndim * [slice(None)]
        l[i] = slice(None, n_2)
        l = tuple(l)

        r = y.ndim * [slice(None)]
        r[i] = slice(n_2, None)
        r = tuple(r)

        y = _concatenate([y[r], y[l]], axis=i)

        if len(x.chunks[i]) == 1:
            y = y.rechunk({i: x.chunks[i]})

    return y

@derived_from(np.fft)
def fftshift(x, axes=None):
    return _fftshift_helper(x, axes=axes, inverse=False)

@derived_from(np.fft)
def ifftshift(x, axes=None):
    return _fftshift_helper(x, axes=axes, inverse=True)


from typing import Optional, List, Union

import numpy as np
from GPy.kern import Kern
from GPy.kern.src.kern import CombinationKernel
from graphviz import Source
from scipy.spatial.distance import cdist, pdist
from sympy import pprint, latex, mathml, dotprint

from src.autoks.backend.kernel import RawKernelType, kernel_to_infix_tokens, tokens_to_str, sort_kernel, additive_form, \
    is_base_kernel, subkernel_expression, kernels_to_kernel_vecs, is_prod_kernel, is_sum_kernel, compute_kernel, \
    KERNEL_DICT, set_priors
from src.autoks.core.hyperprior import HyperpriorMap
from src.autoks.core.kernel_encoding import kernel_to_tree, KernelTree
from src.autoks.symbolic.kernel_symbol import KernelSymbol
from src.autoks.symbolic.util import postfix_tokens_to_symbol
from src.autoks.util import remove_duplicates
from src.evalg.encoding import infix_tokens_to_postfix_tokens
from src.evalg.serialization import Serializable

class Covariance(Serializable):
    """A wrapper for a GPy Kern"""

    def __init__(self, kernel: RawKernelType):
        self.raw_kernel = kernel

    @property
    def raw_kernel(self) -> RawKernelType:
        return self._raw_kernel

    @raw_kernel.setter
    def raw_kernel(self, new_kernel: RawKernelType) -> None:
        if not isinstance(new_kernel, RawKernelType):
            raise TypeError(f'kernel must be {RawKernelType.__name__}. Found type {new_kernel.__class__.__name__}.')
        self._raw_kernel = new_kernel
        # Set other raw_kernel parameters
        self.infix_tokens = kernel_to_infix_tokens(self.raw_kernel)
        self.postfix_tokens = infix_tokens_to_postfix_tokens(self.infix_tokens)
        self.infix = tokens_to_str(self.infix_tokens, show_params=False)
        self.infix_full = tokens_to_str(self.infix_tokens, show_params=True)
        self.postfix = tokens_to_str(self.postfix_tokens, show_params=False)
        postfix_token_symbols = tokens_to_kernel_symbols(self.postfix_tokens)
        self.symbolic_expr = postfix_tokens_to_symbol(postfix_token_symbols)
        self.symbolic_expr_expanded = self.symbolic_expr.expand()

    def to_dict(self) -> dict:
        input_dict = super().to_dict()
        input_dict["kernel"] = self.raw_kernel.to_dict()
        return input_dict

    @classmethod
    def _format_input_dict(cls, input_dict: dict) -> dict:
        input_dict = super()._format_input_dict(input_dict)
        input_dict["kernel"] = Kern.from_dict(input_dict["kernel"])
        return input_dict

    def to_binary_tree(self) -> KernelTree:
        """Get the binary tree representation of the kernel

        :return:
        """
        return kernel_to_tree(self.raw_kernel)

    def canonical(self) -> RawKernelType:
        """Get canonical form of backend kernel.

        :return:
        """
        return sort_kernel(self.raw_kernel)

    def to_additive_form(self) -> RawKernelType:
        """Convert the kernel to additive form.

        :return:
        """
        return additive_form(self.raw_kernel)

    def pretty_print(self) -> None:
        """Pretty print the kernel.

        :return:
        """
        pprint(self.symbolic_expr)

    def print_full(self) -> None:
        """Print the verbose version of the kernel.

        :return:
        """
        print(self.infix_full)

    def is_base(self) -> bool:
        """Determine whether backend kernel is a 1-d base kernel."""
        return is_base_kernel(self.raw_kernel)

    def is_sum(self) -> bool:
        """Determine whether backend kernel is a sum kernel."""
        return is_sum_kernel(self.raw_kernel)

    def is_prod(self) -> bool:
        """Determine whether backend kernel is a product kernel."""
        return is_prod_kernel(self.raw_kernel)

    def priors(self) -> Optional:
        """Get the priors of the kernel."""
        raise NotImplementedError('This will be implemented soon')

    def set_hyperpriors(self, hyperpriors: HyperpriorMap) -> None:
        inv_KERNEL_DICT = {v: k for k, v in KERNEL_DICT.items()}

        def set_kern_prior(x):
            if not isinstance(x, CombinationKernel) and isinstance(x, Kern):
                cls_name = inv_KERNEL_DICT[x.__class__]
                set_priors(x, hyperpriors[cls_name], in_place=True)

        for part in self.infix_tokens:
            set_kern_prior(part)

    def symbolically_equals(self, other) -> bool:
        """Determine whether this covariance's kernel expression is the same as another's kernel expression."""
        return self.symbolic_expr == other.symbolic_expr

    def symbolic_expanded_equals(self, other) -> bool:
        """Determine whether this covariance's expanded kernel expression is the same as another's expanded kernel
        expression."""
        return self.symbolic_expr_expanded == other.symbolic_expr_expanded

    def infix_equals(self, other) -> bool:
        """Determine whether this covariance's kernel infix expression is the same as another's infix kernel
        expression."""
        # naively compare based on infix
        return isinstance(other, Covariance) and other.infix == self.infix

    def as_latex(self) -> str:
        """Get a LaTeX representation of this covariance."""
        return latex(self.symbolic_expr)

    def as_mathml(self) -> str:
        """Get a MathML representation of this covariance."""
        return mathml(self.symbolic_expr)

    def as_dot(self) -> str:
        """Get a DOT representation of this covariance."""
        return dotprint(self.symbolic_expr)

    def as_graph(self) -> Source:
        """Get a GraphViz Source representation of this covariance."""
        return Source(self.as_dot())

    def __add__(self, other):
        return Covariance(self.raw_kernel + other.raw_kernel)

    def __mul__(self, other):
        return Covariance(self.raw_kernel * other.raw_kernel)

    def __str__(self):
        return str(self.symbolic_expr)

    def __repr__(self):
        return f'{self.__class__.__name__}('f'kernel={self.infix_full !r})'

def pretty_print_covariances(covariances: List[Covariance],
                             kernel_type_label: Optional[str] = None):
    """Pretty print a list of covariances."""
    n_kernels = len(covariances)

    plural_suffix = 's' if n_kernels > 1 else ''
    ending = f'kernel{plural_suffix}:'
    if kernel_type_label is not None:
        message = f'{n_kernels} {kernel_type_label} {ending}'
    else:
        message = f'{n_kernels} {ending}'
    message = message.capitalize()
    print(message)
    for cov in covariances:
        cov.pretty_print()
    print('')

# Symbolic interface
def tokens_to_kernel_symbols(tokens: List[Union[str, RawKernelType]]) -> List[Union[str, KernelSymbol]]:
    symbols = []
    for token in tokens:
        if isinstance(token, str):
            symbols.append(token)
        elif isinstance(token, RawKernelType):
            name = subkernel_expression(token)
            symbols.append(KernelSymbol(name, token))
    return symbols

def euclidean_distance(x: np.ndarray,
                       y: np.ndarray) -> float:
    return np.linalg.norm(x - y)

def kernel_l2_dist(kernel_1: RawKernelType,
                   kernel_2: RawKernelType,
                   x: np.ndarray) -> float:
    """Euclidean distance between two kernel matrices.

    :param kernel_1:
    :param kernel_2:
    :param x:
    :return:
    """

    return euclidean_distance(compute_kernel(kernel_1, x), compute_kernel(kernel_2, x))

def covariance_distance(covariances: List[Covariance],
                        x: np.ndarray) -> np.ndarray:
    """Euclidean distance of all pairs gp_models.

    :param covariances:
    :param x:
    :return:
    """
    # For each pair of kernel matrices, compute Euclidean distance
    n_kernels = len(covariances)
    dists = np.zeros((n_kernels, n_kernels))
    for i in range(n_kernels):
        for j in range(i + 1, n_kernels):
            dists[i, j] = kernel_l2_dist(covariances[i].raw_kernel, covariances[j].raw_kernel, x)
    # Make symmetric
    dists = (dists + dists.T) / 2.
    return dists

def remove_duplicate_kernels(covariances: List[Covariance]) -> List[Covariance]:
    """Remove duplicate gp_models.

    :param covariances:
    :return:
    """
    return remove_duplicates([cov.symbolic_expr for cov in covariances], covariances)

def kernel_vec_avg_dist(kvecs1: np.ndarray,
                        kvecs2: np.ndarray) -> float:
    """Average Euclidean distance between two lists of vectors.

    :param kvecs1: n_1 x d array encoding of an additive kernel part
    :param kvecs2: n_2 x d array encoding of an additive kernel part
    :return:
    """
    dists = cdist(kvecs1, kvecs2, metric='euclidean')
    return float(np.mean(dists))

def all_pairs_avg_dist(covariances: List[Covariance],
                       base_kernels: List[str],
                       n_dims: int) -> float:
    """Mean distance between all pairs of gp_models.

    Can be thought of as a diversity score of a population of gp_models
    :param covariances:
    :param base_kernels:
    :param n_dims:
    :return:
    """
    if len(covariances) < 2:
        return 0.

    raw_kernels = [cov.raw_kernel for cov in covariances]
    kernel_vecs = kernels_to_kernel_vecs(raw_kernels, base_kernels, n_dims)

    # compute average Euclidean distance for all pairs of gp_models
    data = np.empty((len(kernel_vecs), 1), dtype=np.object)
    for i, kvec in enumerate(kernel_vecs):
        data[i, 0] = kvec
    pairwise_dist = pdist(data, metric=lambda u, v: kernel_vec_avg_dist(u[0], v[0]))
    return float(np.mean(pairwise_dist))

def inner_frob(m, n):
    """Frobenius inner product"""
    return np.trace(m.T.conjugate() @ n)

def alignment(k1: np.ndarray, k2: np.ndarray) -> float:
    """Alignment A(k1, k2) between two kernel matrices

    It can be viewed as the cosine of the angle between the matrices viewed as 2-d vectors

    0 <= A(k1, k2) <= 1

        Alignment $A$ between two kernel matrices $K_1$ and $K_2$:

    $$A(K_1, K_2) = \frac{\langle K_1, K_2 \rangle_F}{\sqrt{\langle K_1, K_1 \rangle_F \langle K_2, K_2 \rangle_F}}$$
    """
    k1_dot_k2 = inner_frob(k1, k2)
    k1_dot_k1 = inner_frob(k1, k1)
    k2_dot_k2 = inner_frob(k2, k2)

    return k1_dot_k2 / np.sqrt(k1_dot_k1 * k2_dot_k2)

def centered_alignment(k1: np.ndarray, k2: np.ndarray) -> float:
    """Centered kernel alignment

    Cortes et al. (2012)
    """
    k1_centered = center_kernel(k1)
    k2_centered = center_kernel(k2)
    return alignment(k1_centered, k2_centered)

def center_kernel(k: np.ndarray) -> np.ndarray:
    """Center a kernel matrix"""
    m = k.shape[0]
    identity = np.eye(m)
    ones = np.ones((m, 1))
    centering = (identity - (ones @ ones.T) / m)
    return centering @ k @ centering

def pairwise_centered_alignments(covariances: List[Covariance],
                                 x: np.ndarray) -> np.ndarray:
    """Alignment of all pairs of covariances.

    :param covariances:
    :param x:
    :return:
    """
    # For each pair of kernel matrices, compute alignment
    n_kernels = len(covariances)
    dists = np.zeros((n_kernels, n_kernels))
    for i in range(n_kernels):
        for j in range(i + 1, n_kernels):
            k1 = compute_kernel(covariances[i].raw_kernel, x)
            k2 = compute_kernel(covariances[j].raw_kernel, x)
            dists[i, j] = centered_alignment(k1, k2)
    # Make symmetric
    dists = (dists + dists.T) / 2.
    return dists


# Copyright 2018 Timo Nolle
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ==============================================================================

import numpy as np

from april.anomalydetection.utils import label_collapse
from april.anomalydetection.utils import max_collapse
from april.anomalydetection.utils.heuristic import best_heuristic
from april.anomalydetection.utils.heuristic import elbow_heuristic
from april.anomalydetection.utils.heuristic import ratio_heuristic
from april.anomalydetection.utils.heuristic import lowest_plateau_heuristic
from april.enums import Base
from april.enums import Class
from april.enums import Heuristic
from april.enums import Strategy

class Binarizer(object):
    def __init__(self, result, mask, features, targets=None):
        self.result = result
        self._mask = mask
        self.mask_ = mask
        self.features = features
        self._targets = targets

        #Try to fix dimensions
        if self.mask_.shape != self.result.scores.shape:
            if len(self.mask_) != len(self.result.scores.shape):
                self.mask_ = np.expand_dims(self.mask_, axis=-1)
            self.mask_ = np.repeat(self.mask_, self.result.scores.shape[-1], axis=-1)

        self.targets = None
        if self._targets is not None:
            if self.result.scores.shape[2] == 1:
                self._targets = np.delete(self._targets,1, axis=2) # 删除第二列
            elif self.result.scores.shape[2] == 2:
                pass
            self.targets = dict((a, self.mask(label_collapse(self._targets, axis=a))) for a in [0, 1,2])

    def mask(self, a):
        print('**************a.shape=', a.shape)
        if len(a.shape) == 1:
            print('1 len(a.shape)=', len(a.shape))
            m = self.mask_[:, 0, 0]
        elif len(a.shape) == 2:
            print('2 len(a.shape)=', len(a.shape))
            m = self.mask_[:, :, 0]
        else:
            print('? len(a.shape)=', len(a.shape))
            m = self.mask_
        print('type(a)=', type(a), ',type(m)=', type(m))
        print('a.shape=', a.shape, ', m.shape=', m.shape)
#         print('a[0]=', a[0], ', m[0]=', m[0])
        return np.ma.array(a, mask=m)

    def get_targets(self, axis=2):
        return self.targets.get(axis)

    def correct_shape(self, tau, strategy):
        tau = np.asarray(tau)
        if strategy == Strategy.POSITION:
            tau = tau[:, None]
        if strategy == Strategy.POSITION_ATTRIBUTE:
            tau = tau.reshape(*self.result.scores.shape[1:])
        return tau

    def split_by_strategy(self, a, strategy):
        if strategy == Strategy.SINGLE:
            return [a]
        elif isinstance(a, list):
            if strategy == Strategy.POSITION:
                return [[_a[:, i:i + 1] for _a in a] for i in range(len(a[0][0]))]
            elif strategy == Strategy.ATTRIBUTE:
                return [[_a] for _a in a]
            elif strategy == Strategy.POSITION_ATTRIBUTE:
                return [[_a[:, i:i + 1]] for i in range(len(a[0][0])) for _a in a]
        else:
            if strategy == Strategy.POSITION:
                return [a[:, i:i + 1, :] for i in range(a.shape[1])]
            elif strategy == Strategy.ATTRIBUTE:
                return [a[:, :, i:i + 1] for i in range(a.shape[2])]
            elif strategy == Strategy.POSITION_ATTRIBUTE:
                return [a[:, i:i + 1, j:j + 1] for i in range(a.shape[1]) for j in range(a.shape[2])]

    def get_grid_candidate_taus(self, a, steps=20, axis=0):
        """G in the paper."""
        return np.linspace(max_collapse(a, axis=axis).min() - .001, a.max(), steps)

    def get_candidate_taus(self, a, axis=0):
        a = max_collapse(a, axis=axis).compressed()
        a_min = a.min()
        a_max = a.max()
        if a_max > a_min:
            a = (a_max - a) / (a_max - a_min)
        a = 2 * (a / 2).round(2)
        if a_max > a_min:
            a = a_max - a * (a_max - a_min)
        a = np.sort(np.unique(a))
        a[0] -= .001
        if len(a) < 5:
            a = np.linspace(a_min - .001, a_max, 5)
        return a

    def get_legacy_tau(self, scores, heuristic=Heuristic.DEFAULT, strategy=Strategy.SINGLE, axis=0):
        if heuristic == Heuristic.DEFAULT:
            return np.array([0.5])

        if not isinstance(scores, np.ma.MaskedArray):
            scores = self.mask(scores)

        alpha = None
        if strategy == Strategy.SINGLE:
            alpha = np.array([scores.mean()])
        elif strategy == Strategy.ATTRIBUTE:
            alpha = scores.mean(axis=1).mean(axis=0).data
        elif strategy == Strategy.POSITION:
            alpha = scores.mean(axis=2).mean(axis=0).data[:, None]
        elif strategy == Strategy.POSITION_ATTRIBUTE:
            alpha = scores.mean(axis=0).data

        taus = self.get_grid_candidate_taus(scores / alpha, axis=axis)
        tau = None
        if heuristic == Heuristic.BEST:
            y_true = self.get_targets(axis=axis)
            tau = best_heuristic(taus=taus, theta=self.legacy_binarize, y_true=y_true, alpha=alpha, scores=scores,
                                 axis=axis)

        if heuristic == Heuristic.RATIO:
            tau = ratio_heuristic(taus=taus, theta=self.legacy_binarize, scores=scores, axis=axis, alpha=alpha)

        if heuristic in [Heuristic.ELBOW_DOWN, Heuristic.ELBOW_UP]:
            tau = elbow_heuristic(taus=taus, theta=self.legacy_binarize, scores=scores, axis=axis,
                                  alpha=alpha)[heuristic]

        if heuristic in [Heuristic.LP_LEFT, Heuristic.LP_MEAN, Heuristic.LP_RIGHT]:
            tau = lowest_plateau_heuristic(taus=taus, theta=self.legacy_binarize, scores=scores, axis=axis,
                                           alpha=alpha)[heuristic]

        return tau * alpha

    def get_tau(self, scores, heuristic=Heuristic.DEFAULT, strategy=Strategy.SINGLE, axis=0, taus=None):
        if heuristic == Heuristic.DEFAULT:
            return np.array([0.5])

        if not isinstance(scores, np.ma.MaskedArray):
            scores = self.mask(scores)

        scores = self.split_by_strategy(scores, strategy)

        if heuristic in [Heuristic.MEAN, Heuristic.MEDIAN]:
            scores = [max_collapse(s, axis=axis) for s in scores]
            if heuristic == Heuristic.MEAN:
                return self.correct_shape([np.mean(s[np.round(s, 1) > 0]) for s in scores], strategy)
            elif heuristic == Heuristic.MEDIAN:
                return self.correct_shape([np.median(s[np.round(s, 1) > 0]) for s in scores], strategy)

        if taus is None:
            taus = [self.get_candidate_taus(s, axis=axis) for s in scores]
        else:
            taus = [taus] * len(scores)

        tau = None
        if heuristic == Heuristic.BEST:
            y_trues = self.split_by_strategy(self.get_targets(axis=2), strategy)
            y_trues = [label_collapse(y, axis=axis) for y in y_trues]
            tau = [best_heuristic(taus=t, theta=self.threshold_binarize, y_true=y, scores=s, axis=axis)
                   for s, t, y in zip(scores, taus, y_trues)]

        if heuristic == Heuristic.RATIO:
            tau = [ratio_heuristic(taus=t, scores=s, theta=self.threshold_binarize, axis=axis)
                   for s, t in zip(scores, taus)]

        if heuristic in [Heuristic.ELBOW_DOWN, Heuristic.ELBOW_UP]:
            tau = [elbow_heuristic(taus=t, scores=s, theta=self.threshold_binarize, axis=axis)[heuristic]
                   for s, t in zip(scores, taus)]

        if heuristic in [Heuristic.LP_LEFT, Heuristic.LP_MEAN, Heuristic.LP_RIGHT]:
            tau = [lowest_plateau_heuristic(taus=t, scores=s, theta=self.threshold_binarize, axis=axis)[heuristic]
                   for s, t in zip(scores, taus)]

        return self.correct_shape(tau, strategy)

    def legacy_binarize(self, scores, tau, alpha, axis=0):
        # Apply the threshold function (Theta in the paper) using alpha as a scaling factor
        return self.threshold_binarize(tau=tau * alpha, scores=scores, axis=axis)

    def threshold_binarize(self, tau, scores, axis=0):
        # Apply the threshold function (Theta in the paper)
        predictions = np.array(scores.data > tau, dtype=int)

        # Apply mask
        predictions = np.ma.array(predictions, mask=scores.mask)

        # Positive axis flatten predictions
        if axis in [0, 1]:
            predictions = label_collapse(predictions, axis=axis)

        return predictions

    def binarize(self, scores=None, tau=None, base=None, heuristic=None, strategy=None, go_backwards=False,
                 return_parameters=False, axis=2, heuristic_axis=None):

        if heuristic_axis is None:
            heuristic_axis = axis

        if scores is None:
            if go_backwards:
                scores = self.result.scores_backward
            else:
                scores = self.result.scores

        if not isinstance(scores, np.ma.MaskedArray):
            scores = self.mask(scores)

        # Get baseline threshold (tau in the paper)
        if tau is None or heuristic != Heuristic.MANUAL:
            if base == Base.LEGACY:
                tau = self.get_legacy_tau(scores=scores, heuristic=heuristic, strategy=strategy, axis=heuristic_axis)
            else:
                tau = self.get_tau(scores=scores, heuristic=heuristic, strategy=strategy, axis=heuristic_axis)

        # Apply the threshold function (Theta in the paper)
        predictions = self.threshold_binarize(scores=scores, tau=tau, axis=axis)

        if return_parameters:
            return predictions, tau

        return predictions

    @staticmethod
    def get_scores(probabilities):
        scores = np.zeros_like(probabilities)
        for i in range(scores.shape[2]):
            p = probabilities[:, :, i:i + 1]
            _p = np.copy(probabilities)
            _p[_p <= p] = 0
            scores[:, :, i] = _p.sum(axis=2)
        return scores

    def classify(self, tau, features, predictions):
        def mask(a, mask):
            b = np.copy(a)
            b[mask == 1] = 0
            c = np.copy(a)
            c[mask == 0] = 0
            return b, c

        classification = np.zeros_like(predictions)
        c_cf = classification[:, :, 0]
        c_data = classification[:, :, 1:]
        predictions_cf = predictions[:, :, 0]
        predictions_data = predictions[:, :, 1:]

        # Attribute heuristic
        c_data[predictions_data == 1] = Class.ATTRIBUTE
        c_data[predictions_data == 0] = Class.NORMAL_ATTRIBUTE

        # Insert and Skip heuristics
        if self.result.predictions is not None:
            f = features[0]

            # Top-1 predictions
            # p = np.argmax(self.result.predictions[0], axis=2) + 1

            # Top-n predictions according to threshold (tau)
            _p = self.get_scores(self.result.predictions[0])
            p = np.zeros_like(_p) + np.arange(_p.shape[-1]) + 1
            p[_p > tau[0]] = -1

            # Mask padding
            p[self._mask] = -1

            # Helper objects
            pfht = np.zeros_like(predictions_cf)
            pfhf = np.zeros_like(predictions_cf)
            pftt = np.zeros_like(predictions_cf)
            pftf = np.zeros_like(predictions_cf)
            ppht = np.zeros_like(predictions_cf)
            pphf = np.zeros_like(predictions_cf)
            pptt = np.zeros_like(predictions_cf)
            pptf = np.zeros_like(predictions_cf)
            ffht = np.zeros_like(predictions_cf)
            ffhf = np.zeros_like(predictions_cf)
            fftt = np.zeros_like(predictions_cf)
            fftf = np.zeros_like(predictions_cf)
            fpht = np.zeros_like(predictions_cf)
            fphf = np.zeros_like(predictions_cf)
            fptt = np.zeros_like(predictions_cf)
            fptf = np.zeros_like(predictions_cf)

            for j in np.arange(f.shape[1]):
                # Current prediction and feature
                _p = p[:, j:j + 1]
                _f = f[:, j:j + 1]

                # Top-1 Predictions
                ph = p[:, :j]
                phf, pht = mask(ph, predictions_cf[:, :j])
                pt = p[:, j + 1:]
                ptf, ptt = mask(pt, predictions_cf[:, j + 1:])

                # Actual case features
                fh = f[:, :j]
                fhf, fht = mask(fh, predictions_cf[:, :j])
                ft = f[:, j + 1:]
                ftf, ftt = mask(ft, predictions_cf[:, j + 1:])

                # Prediction appears elsewhere in case
                pfht[:, j] = np.any(np.any(_p == fht[:, :, np.newaxis], axis=-1), axis=-1)
                pfhf[:, j] = np.any(np.any(_p == fhf[:, :, np.newaxis], axis=-1), axis=-1)
                pftt[:, j] = np.any(np.any(_p == ftt[:, :, np.newaxis], axis=-1), axis=-1)
                pftf[:, j] = np.any(np.any(_p == ftf[:, :, np.newaxis], axis=-1), axis=-1)

                # Prediction appears elsewhere in predictions
                ppht[:, j] = np.any(np.any(_p == pht, axis=-1), axis=-1)
                pphf[:, j] = np.any(np.any(_p == phf, axis=-1), axis=-1)
                pptt[:, j] = np.any(np.any(_p == ptt, axis=-1), axis=-1)
                pptf[:, j] = np.any(np.any(_p == ptf, axis=-1), axis=-1)

                # Event appears elsewhere in case
                ffht[:, j] = np.any(_f == fht, axis=-1)
                ffhf[:, j] = np.any(_f == fhf, axis=-1)
                fftt[:, j] = np.any(_f == ftt, axis=-1)
                fftf[:, j] = np.any(_f == ftf, axis=-1)

                # Event appears elsewhere in predictions
                fpht[:, j] = np.any(np.any(_f[:, :, np.newaxis] == pht, axis=-1), axis=-1)
                fphf[:, j] = np.any(np.any(_f[:, :, np.newaxis] == phf, axis=-1), axis=-1)
                fptt[:, j] = np.any(np.any(_f[:, :, np.newaxis] == ptt, axis=-1), axis=-1)
                fptf[:, j] = np.any(np.any(_f[:, :, np.newaxis] == ptf, axis=-1), axis=-1)

            # Classification rules
            skips = np.logical_and(predictions_cf == 1, ~np.logical_or(pfhf, pftf))
            inserts = np.logical_and(predictions_cf == 1, np.logical_or(pfhf, pftf))
            reworks = np.logical_and(predictions_cf == 1, ffhf)

            shifts = np.logical_and(predictions_cf == 1, np.logical_xor(pfht, pftt))
            lates = np.logical_and(predictions_cf == 1, fpht)
            earlies = np.logical_and(predictions_cf == 1, fptt)

            # Set the labels
            c_cf[inserts] = Class.INSERT
            c_cf[skips] = Class.SKIP
            c_cf[shifts] = Class.SHIFT
            c_cf[lates] = Class.LATE
            c_cf[earlies] = Class.EARLY
            c_cf[reworks] = Class.REWORK

        return self.mask(classification)


"""

Filename: utils.py

Purpose: Set of utility functions

Author: John Weaver
Date created: 28.11.2018
Possible problems:
1.

"""
import os
import numpy as np
from tractor.galaxy import ExpGalaxy
from tractor import EllipseE
from tractor.galaxy import ExpGalaxy
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm, SymLogNorm
from matplotlib.patches import Ellipse
from matplotlib.patches import Rectangle
from skimage.segmentation import find_boundaries

import config as conf
import matplotlib.cm as cm
import random
from time import time
from astropy.io import fits

import logging
logger = logging.getLogger('farmer.utils')

def header_from_dict(params):
    """ Take in dictionary and churn out a header. Never forget configs again. """
    hdr = fits.Header()
    total_public_entries = np.sum([ not k.startswith('__') for k in params.keys()])
    logger.debug(f'header_from_dict :: Dictionary has {total_public_entries} entires')
    tstart = time()
    for i, attr in enumerate(params.keys()):
        if not attr.startswith('__'):
            logger.debug(f'header_from_dict ::   {attr}')
            value = params[attr]
            if type(value) == str:
                # store normally
                hdr.set(f'CONF{i+1}', value, attr)
            if type(value) in (float, int):
                # store normally
                hdr.set(f'CONF{i+1}', value, attr)
            if type(value) in (list, tuple):
                # freak out.
                for j, val in enumerate(value):
                    hdr.set(f'CONF{i+1}_{j+1}', str(val), f'{attr}_{j+1}')

    logger.debug(f'header_from_dict :: Completed writing header ({time() - tstart:2.3f}s)')
    return hdr

def create_circular_mask(h, w, center=None, radius=None):

    if center is None: # use the middle of the image
        center = [int(w/2), int(h/2)]
    if radius is None: # use the smallest distance between the center and image walls
        radius = min(center[0], center[1], w-center[0], h-center[1])

    Y, X = np.ogrid[:h, :w]
    dist_from_center = np.sqrt((X - center[0])**2 + (Y-center[1])**2)

    mask = np.zeros((h, w), dtype=int)
    mask[dist_from_center <= radius] = 1
    return mask

class SimpleGalaxy(ExpGalaxy):
    '''This defines the 'SIMP' galaxy profile -- an exponential profile
    with a fixed shape of a 0.45 arcsec effective radius and spherical
    shape.  It is used to detect marginally-resolved galaxies.
    '''
    shape = EllipseE(0.45 / conf.PIXEL_SCALE, 0., 0.)

    def __init__(self, *args):
        super(SimpleGalaxy, self).__init__(*args)

    def __str__(self):
        return (self.name + ' at ' + str(self.pos)
                + ' with ' + str(self.brightness))

    def __repr__(self):
        return (self.name + '(pos=' + repr(self.pos) +
                ', brightness=' + repr(self.brightness) + ')')

    @staticmethod
    def getNamedParams():
        return dict(pos=0, brightness=1)

    def getName(self):
        return 'SimpleGalaxy'

    ### HACK -- for Galaxy.getParamDerivatives()
    def isParamFrozen(self, pname):
        if pname == 'shape':
            return True
        return super(SimpleGalaxy, self).isParamFrozen(pname)

def make_weights(fn):

    # Grab an image, estimate rms, make weight map fits file

    import sys
    from astropy.stats import sigma_clipped_stats

    hdul = fits.open(fn)

    print(hdul.info())

    __, __, rms = sigma_clipped_stats(hdul[0].data)

    wgt = np.zeros_like(hdul[0].data)
    wgt[rms>0] = 1/rms**2

    hdul[0].data = wgt

    p = fn.split('.fits')
    fout = p[0]+conf.WEIGHT_EXT+'.fits'
    print('Writing weight map to ', fout)
    hdul.writeto(fout)


# AUTOGENERATED! DO NOT EDIT! File to edit: data_vis.ipynb (unless otherwise specified).

__all__ = ['get_settings', 'cross_section_helper']

# Cell
import numpy as np
import scipy.interpolate
import matplotlib.pyplot as plt

# Cell
def get_settings():
    """
    Returns some nice default settings for matplotlib to be used with `matplotlib.pyplot.rc_context`.
    """
    return {'axes.labelsize': 32,
            'xtick.major.size': 10,
            'xtick.major.width': 1.5,
            'xtick.labelsize': 24,
            'ytick.major.size': 10,
            'ytick.major.width': 1.5,
            'ytick.labelsize': 24,
            'legend.fontsize': 18,
            'lines.linewidth': 4,
            'lines.markersize': 10,
            'figure.figsize': (12,8)}

# Cell
class cross_section_helper:
    """
    Class enabling computations and visualizations relating to signal and background cross sections
    """

    def __init__(self, masses, sig_css, bg_css, mass_units='GeV'):
        self.masses = masses
        self.sig_css = sig_css
        self.bg_css = bg_css
        self.mass_units = mass_units
        self.log_sig_css = np.log10(self.sig_css)
        self.m2logcsF = scipy.interpolate.interp1d(self.masses, self.log_sig_css, kind='cubic')
        self.logcs2mF = scipy.interpolate.interp1d(self.m2logcsF(self.masses), self.masses, kind='cubic')
        # WHERE TO PUT THIS VARIABLE?
        self.conv = 10**15 / 10**12 # cross sec (pb) * lumi (fb^{-1}) * self.conv = number of events

    def sig_cs(self, masses, who='E'):
        r"""
        Given a mass, returns the signal cross section (in pb) through interpolation
        """
        return np.power(10, self.m2logcsF(masses))

    def mass(self, css, who='E'):
        r"""
        Given a cross section, returns the particle mass value yielding that signal cross section
        """
        return self.logcs2mF(np.log10(css))

    def absolute_max_mass_sens(self, lumi=3000, sig=5):
        """
        Given a luminosity and desired signal significance, returns the best-case scenario highest probeable mass
        (i.e., TPR = 1, FPR = 0)
        """
        return self.mass(sig**2 / (lumi * self.conv))

    def max_mass_sens_versus_tpr_fpr(self, plot=True, sig=5, lumi=3000, tpr_bounds=(0.1,1), fpr_bounds=(10**-7, 0.5), res=1000,
                                     cvalues=None, clabels=None, manual=None):
        """
        For fixed luminosity and desired signal significance, plot the maximum mass sensitivity as a function of a
        binary classifier's true positive rate and false positive rate.
        """
        tprs = np.linspace(tpr_bounds[0], tpr_bounds[1], 100)
        fprs = np.sort(np.logspace(np.log10(fpr_bounds[0]), np.log10(fpr_bounds[1]), num=100, base=10))
        bg_yield = self.conv * np.sum(self.bg_css) * lumi
        coefsss = [
            [[-(self.conv * tpr * lumi)**2, sig**2 * self.conv * tpr * lumi, sig**2 * fpr * bg_yield] for fpr in fprs] for tpr in tprs]
        sig_csss = [[np.amax(np.roots(coefs)) for coefs in coefss] for coefss in coefsss]
        max_mass_senss = self.mass(sig_csss)

        # have found this plot works best when plotting log(fpr) on y-axis
        log_fprs = np.log10(fprs)
        ext = [tprs[0], tprs[-1], log_fprs[0], log_fprs[-1]]
        pts = np.array([[tpr,log_fpr] for tpr in tprs for log_fpr in log_fprs])
        data = max_mass_senss.flatten()
        grid = np.array(
                [[[x,y] for x in np.linspace(ext[0], ext[1], res)]
                for y in np.linspace(ext[2], ext[3], res)])
        interp = scipy.interpolate.griddata(pts, data, grid)

        if plot:
            with plt.rc_context(get_settings()):
                plt.imshow(interp, origin='lower', aspect='auto', extent=ext, cmap='winter')
                cbar = plt.colorbar()
                plt.xlabel(r'{}'.format('True Positive Rate'))
                plt.ylabel(r'{}'.format('False Postive Rate'))
                cbar.set_label(r'{}'.format(f"Mass Sensitivity ({self.mass_units})"), rotation=90)
                if None not in [cvalues, clabels]:
                    result = plt.contour(grid[:,:,0], grid[:,:,1], interp, cvalues, colors='white',
                                                 linewidths=3, linestyles='dashed')
                    fmt = {lev:lab for lev, lab in zip(result.levels, clabels)}
                    plt.clabel(result, result.levels, inline=True, fmt=fmt, fontsize=32, manual=manual)
                minytick, maxytick = [int(log_fprs[0]) + 1, int(log_fprs[-1]) - 1]
                plt.yticks(
                    [i for i in range(minytick, maxytick+1)], [r'$10^{' + f'{i}' + r'}$' for i in range(minytick, maxytick+1)])

        return [grid, interp]