# -*- mode: python; coding: utf-8 -*-
# Copyright (c) 2019 Radio Astronomy Software Group
# Licensed under the 2-clause BSD License

"""Class for reading MWA correlator FITS files."""
import warnings
import itertools
import numpy as np
from astropy.io import fits
from astropy.time import Time
from astropy import constants as const
from pyuvdata.data import DATA_PATH

from .. import _corr_fits

from . import UVData
from .. import utils as uvutils

__all__ = ["input_output_mapping", "MWACorrFITS"]

def input_output_mapping():
    """Build a mapping dictionary from pfb input to output numbers."""
    # the polyphase filter bank maps inputs to outputs, which the MWA
    # correlator then records as the antenna indices.
    # the following is taken from mwa_build_lfiles/mwac_utils.c
    # inputs are mapped to outputs via pfb_mapper as follows
    # (from mwa_build_lfiles/antenna_mapping.h):
    # floor(index/4) + index%4 * 16 = input
    # for the first 64 outputs, pfb_mapper[output] = input
    return _corr_fits.input_output_mapping()

class MWACorrFITS(UVData):
    """
    UVData subclass for reading MWA correlator fits files.

    This class should not be interacted with directly; instead use the
    read_mwa_corr_fits method on the UVData class.
    """

    def correct_cable_length(self, cable_lens):
        """
        Apply a cable length correction to the data array.

        Parameters
        ----------
        cable_lens : list of strings
        A list of strings containing the cable lengths for each antenna.
        """
        # from MWA_Tools/CONV2UVFITS/convutils.h
        cable_len_diffs = _corr_fits.get_cable_len_diffs(
            self.Nblts, self.ant_1_array, self.ant_2_array, cable_lens
        )
        self.data_array *= np.exp(
            -1j
            * 2
            * np.pi
            * cable_len_diffs.reshape(self.Nblts, 1)
            / const.c.to("m/s").value
            * self.freq_array.reshape(1, self.Nfreqs)
        )[:, :, None]

    def flag_init(
        self,
        num_fine_chan,
        edge_width=80e3,
        start_flag=2.0,
        end_flag=0.0,
        flag_dc_offset=True,
    ):
        """
        Apply routine flagging to the MWA Correlator FITS file data.

        Includes options to flag the coarse channel edges, beginning and end
        of obs, as well as the center fine channel of each coarse channel.

        Parameters
        ----------
        edge_width: float
            The width to flag on the edge of each coarse channel, in hz. Set to
            0 for no edge flagging.
        start_flag: float
            The number of seconds to flag at the beginning of the observation.
            Set to 0 for no flagging.
        end_flag: floats
            The number of seconds to flag at the end of the observation. Set to
            0 for no flagging.
        flag_dc_offset: bool
            Set to True to flag the center fine channel of each coarse channel.

        Raises
        ------
        ValueError
            If edge_width is not an integer multiple of the channel_width of
            the data (0 also acceptable).
            If start_flag is not an integer multiple of the integration time
            (0 also acceptable).
            If end_flag is not an integer multiple of the integration time
            (0 also acceptable).

        """
        if (edge_width % self.channel_width) > 0:
            raise ValueError(
                "The edge_width must be an integer multiple of the "
                "channel_width of the data or zero."
            )
        if (start_flag % self.integration_time[0]) > 0:
            raise ValueError(
                "The start_flag must be an integer multiple of the "
                "integration_time of the data or zero."
            )
        if (end_flag % self.integration_time[0]) > 0:
            raise ValueError(
                "The end_flag must be an integer multiple of the "
                "integration_time of the data or zero."
            )

        num_ch_flag = int(edge_width / self.channel_width)
        num_start_flag = int(start_flag / self.integration_time[0])
        num_end_flag = int(end_flag / self.integration_time[0])

        if num_ch_flag > 0:
            edge_inds = []
            for ch_count in range(num_ch_flag):
                # count up from the left
                left_chans = list(range(ch_count, self.Nfreqs, num_fine_chan))
                # count down from the right
                right_chans = list(range(self.Nfreqs - 1 - ch_count, 0, -num_fine_chan))
                edge_inds = edge_inds + left_chans + right_chans

            self.flag_array[:, :, edge_inds, :] = True

        if flag_dc_offset:
            center_inds = list(range(num_fine_chan // 2, self.Nfreqs, num_fine_chan))

            self.flag_array[:, :, center_inds, :] = True

        if (num_start_flag > 0) or (num_end_flag > 0):
            shape = self.flag_array.shape
            # TODO: Spw axis to be collapsed in future release
            # Asserting this here because this is effectively a stripped down UVFITS
            # reader, and thus assuming that this should only support simple tables
            assert shape[1] == 1
            reshape = [self.Ntimes, self.Nbls, 1, self.Nfreqs, self.Npols]
            self.flag_array = np.reshape(self.flag_array, reshape)
            if num_start_flag > 0:
                self.flag_array[:num_start_flag, :, :, :, :] = True
            if num_end_flag > 0:
                self.flag_array[-num_end_flag:, :, :, :, :] = True
            self.flag_array = np.reshape(self.flag_array, shape)

    def read_mwa_corr_fits(
        self,
        filelist,
        use_cotter_flags=None,
        remove_dig_gains=True,
        remove_coarse_band=True,
        correct_cable_len=False,
        phase_to_pointing_center=False,
        propagate_coarse_flags=True,
        flag_init=True,
        edge_width=80e3,
        start_flag="goodtime",
        end_flag=0.0,
        flag_dc_offset=True,
        background_lsts=True,
        read_data=True,
        data_array_dtype=np.complex64,
        nsample_array_dtype=np.float32,
        run_check=True,
        check_extra=True,
        run_check_acceptability=True,
        strict_uvw_antpos_check=False,
    ):
        """
        Read in MWA correlator gpu box files.

        The default settings remove some of the instrumental effects in the bandpass
        by dividing out the digital gains and the coarse band shape.
        If the desired output is raw correlator data, set remove_dig_gains=False,
        remove_coarse_band=False, correct_cable_len=False, and
        phase_to_pointing_center=False.

        Parameters
        ----------
        filelist : list of str
            The list of MWA correlator files to read from. Must include at
            least one fits file and only one metafits file per data set.
            Can also be a list of lists to read multiple data sets.
        axis : str
            Axis to concatenate files along. This enables fast concatenation
            along the specified axis without the normal checking that all other
            metadata agrees. This method does not guarantee correct resulting
            objects. Please see the docstring for fast_concat for details.
            Allowed values are: 'blt', 'freq', 'polarization'. Only used if
            multiple files are passed.
        use_cotter_flags : bool
            Option to use cotter output mwaf flag files. Defaults to true if
            cotter flag files are submitted.
        remove_dig_gains : bool
            Option to divide out digital gains.
        remove_coarse_band : bool
            Option to divide out coarse band shape.
        correct_cable_len : bool
            Option to apply a cable delay correction.
        phase_to_pointing_center : bool
            Option to phase to the observation pointing center.
        propagate_coarse_flags : bool
            Option to propagate flags for missing coarse channel integrations
            across frequency.
        flag_init: bool
            Set to True in order to do routine flagging of coarse channel edges,
            start or end integrations, or the center fine channel of each coarse
            channel. See associated keywords.
        edge_width: float
            Only used if flag_init is True. The width to flag on the edge of
            each coarse channel, in hz. Errors if not equal to integer multiple
            of channel_width. Set to 0 for no edge flagging.
        start_flag: float or str
            Only used if flag_init is True. The number of seconds to flag at the
            beginning of the observation. Set to 0 for no flagging. Default is
            'goodtime', which uses information in the metafits file to determine
            the length of time that should be flagged. Errors if input is not a
            float or 'goodtime'. Errors if float input is not equal to an
            integer multiple of the integration time.
        end_flag: floats
            Only used if flag_init is True. The number of seconds to flag at the
            end of the observation. Set to 0 for no flagging. Errors if not
            equal to an integer multiple of the integration time.
        flag_dc_offset: bool
            Only used if flag_init is True. Set to True to flag the center fine
            channel of each coarse channel.
        background_lsts : bool
            When set to True, the lst_array is calculated in a background thread.
        read_data : bool
            Read in the visibility, nsample and flag data. If set to False, only
            the metadata will be read in. Setting read_data to False results in
            a metadata only object.
        data_array_dtype : numpy dtype
            Datatype to store the output data_array as. Must be either
            np.complex64 (single-precision real and imaginary) or np.complex128
            (double-precision real and imaginary).
        nsample_array_dtype : numpy dtype
            Datatype to store the output nsample_array as. Must be either
            np.float64 (double-precision), np.float32 (single-precision), or
            np.float16 (half-precision). Half-precision is only recommended for
            cases where no sampling or averaging of baselines will occur,
            because round-off errors can be quite large (~1e-3).
        run_check : bool
            Option to check for the existence and proper shapes of parameters
            after after reading in the file (the default is True,
            meaning the check will be run).
        check_extra : bool
            Option to check optional parameters as well as required ones (the
            default is True, meaning the optional parameters will be checked).
        run_check_acceptability : bool
            Option to check acceptable range of the values of parameters after
            reading in the file (the default is True, meaning the acceptable
            range check will be done).
        strict_uvw_antpos_check : bool
            Option to raise an error rather than a warning if the check that
            uvws match antenna positions does not pass.

        Raises
        ------
        ValueError
            If required files are missing or multiple files metafits files are
            included in filelist.
            If files from different observations are included in filelist.
            If files in fileslist have different fine channel widths
            If file types other than fits, metafits, and mwaf files are included
            in filelist.

        """
        metafits_file = None
        ppds_file = None
        obs_id = None
        file_dict = {}
        start_time = 0.0
        end_time = 0.0
        included_file_nums = []
        included_flag_nums = []
        cotter_warning = False
        num_fine_chans = 0

        # do datatype checks
        if data_array_dtype not in (np.complex64, np.complex128):
            raise ValueError("data_array_dtype must be np.complex64 or np.complex128")
        if nsample_array_dtype not in (np.float64, np.float32, np.float16):
            raise ValueError(
                "nsample_array_dtype must be one of: np.float64, np.float32, np.float16"
            )
        # do start_flag check
        if not isinstance(start_flag, (int, float)):
            if start_flag != "goodtime":
                raise ValueError("start_flag must be int or float or 'goodtime'")

        # iterate through files and organize
        # create a list of included coarse channels
        # find the first and last times that have data
        for file in filelist:
            if file.lower().endswith(".metafits"):
                # force only one metafits file
                if metafits_file is not None:
                    raise ValueError("multiple metafits files in filelist")
                metafits_file = file
            elif file.lower().endswith(".fits"):
                # check if ppds file
                try:
                    fits.getheader(file, extname="ppds")
                    ppds_file = file
                except Exception:
                    # check obsid
                    head0 = fits.getheader(file, 0)
                    if obs_id is None:
                        obs_id = head0["OBSID"]
                    else:
                        if head0["OBSID"] != obs_id:
                            raise ValueError(
                                "files from different observations submitted "
                                "in same list"
                            )
                    # check headers for first and last times containing data
                    headstart = fits.getheader(file, 1)
                    headfin = fits.getheader(file, -1)
                    first_time = headstart["TIME"] + headstart["MILLITIM"] / 1000.0
                    last_time = headfin["TIME"] + headfin["MILLITIM"] / 1000.0
                    if start_time == 0.0:
                        start_time = first_time
                    # check that files with a timing offset can be aligned
                    elif np.abs(start_time - first_time) % headstart["INTTIME"] != 0.0:
                        raise ValueError(
                            "coarse channel start times are misaligned by an amount that is not \
                                an integer multiple of the integration time"
                        )
                    elif start_time > first_time:
                        start_time = first_time
                    if end_time < last_time:
                        end_time = last_time
                    # get number of fine channels
                    if num_fine_chans == 0:
                        num_fine_chans = headstart["NAXIS2"]
                    elif num_fine_chans != headstart["NAXIS2"]:
                        raise ValueError(
                            "files submitted have different fine channel widths"
                        )
                    # get the file number from the file name;
                    # this will later be mapped to a coarse channel
                    file_num = int(file.split("_")[-2][-2:])
                    if file_num not in included_file_nums:
                        included_file_nums.append(file_num)
                    # organize files
                    if "data" not in file_dict.keys():
                        file_dict["data"] = [file]
                    else:
                        file_dict["data"].append(file)
            # look for flag files
            elif file.lower().endswith(".mwaf"):
                if use_cotter_flags is None:
                    use_cotter_flags = True
                flag_num = int(file.split("_")[-1][0:2])
                included_flag_nums.append(flag_num)
                if use_cotter_flags is False and cotter_warning is False:
                    warnings.warn("mwaf files submitted with use_cotter_flags=False")
                    cotter_warning = True
                elif "flags" not in file_dict.keys():
                    file_dict["flags"] = [file]
                else:
                    file_dict["flags"].append(file)
            else:
                raise ValueError("only fits, metafits, and mwaf files supported")

        # checks:
        if metafits_file is None and ppds_file is None:
            raise ValueError("no metafits file submitted")
        elif metafits_file is None:
            metafits_file = ppds_file
        elif ppds_file is not None:
            ppds = fits.getheader(ppds_file, 0)
            meta = fits.getheader(metafits_file, 0)
            for key in ppds.keys():
                if key not in meta.keys():
                    self.extra_keywords[key] = ppds[key]
        if "data" not in file_dict.keys():
            raise ValueError("no data files submitted")
        if "flags" not in file_dict.keys() and use_cotter_flags:
            raise ValueError(
                "no flag files submitted. Rerun with flag files \
                             or use_cotter_flags=False"
            )

        # first set parameters that are always true
        self.Nspws = 1
        self.spw_array = np.array([0])
        self.phase_type = "drift"
        self.vis_units = "uncalib"
        self.Npols = 4
        self.xorientation = "east"

        # get information from metafits file
        with fits.open(metafits_file, memmap=True) as meta:
            meta_hdr = meta[0].header

            # get a list of coarse channels
            coarse_chans = meta_hdr["CHANNELS"].split(",")
            coarse_chans = np.array(sorted(int(i) for i in coarse_chans))

            # integration time in seconds
            int_time = meta_hdr["INTTIME"]

            # pointing center in degrees
            ra_deg = meta_hdr["RA"]
            dec_deg = meta_hdr["DEC"]
            ra_rad = np.pi * ra_deg / 180
            dec_rad = np.pi * dec_deg / 180

            # set start_flag with goodtime
            if flag_init and start_flag == "goodtime":
                # ppds file does not contain this key
                try:
                    if meta_hdr["GOODTIME"] > start_time:
                        start_flag = meta_hdr["GOODTIME"] - start_time
                    else:
                        start_flag = 0.0
                except KeyError:
                    raise ValueError(
                        "To use start_flag='goodtime', a .metafits file must \
                            be submitted"
                    )

            # get parameters from header
            # this assumes no averaging by this code so will need to be updated
            self.channel_width = float(meta_hdr.pop("FINECHAN") * 1000)
            if "HISTORY" in meta_hdr:
                self.history = str(meta_hdr["HISTORY"])
                meta_hdr.remove("HISTORY", remove_all=True)
            else:
                self.history = ""
            if not uvutils._check_history_version(
                self.history, self.pyuvdata_version_str
            ):
                self.history += self.pyuvdata_version_str
            self.instrument = meta_hdr["TELESCOP"]
            self.telescope_name = meta_hdr.pop("TELESCOP")
            self.object_name = meta_hdr.pop("FILENAME")

            # get rid of the instrument keyword so it doesn't get put back in
            meta_hdr.remove("INSTRUME")
            # get rid of keywords that uvfits.py gets rid of
            bad_keys = ["SIMPLE", "EXTEND", "BITPIX", "NAXIS", "DATE-OBS"]
            for key in bad_keys:
                meta_hdr.remove(key, remove_all=True)
            # store remaining keys in extra keywords
            for key in meta_hdr:
                if key == "COMMENT":
                    self.extra_keywords[key] = str(meta_hdr.get(key))
                elif key != "":
                    self.extra_keywords[key] = meta_hdr.get(key)
            # get antenna data from metafits file table
            meta_tbl = meta[1].data

            # because of polarization, each antenna # is listed twice
            antenna_numbers = meta_tbl["Antenna"][1::2]
            antenna_names = meta_tbl["TileName"][1::2]
            antenna_flags = meta_tbl["Flag"][1::2]
            cable_lens = np.asarray(meta_tbl["Length"][1::2]).astype(np.str_)
            dig_gains = meta_tbl["Gains"][1::2, :].astype(np.float64)

            # get antenna postions in enu coordinates
            antenna_positions = np.zeros((len(antenna_numbers), 3))
            antenna_positions[:, 0] = meta_tbl["East"][1::2]
            antenna_positions[:, 1] = meta_tbl["North"][1::2]
            antenna_positions[:, 2] = meta_tbl["Height"][1::2]

        # reorder antenna parameters from metafits ordering
        reordered_inds = antenna_numbers.argsort()
        self.antenna_numbers = antenna_numbers[reordered_inds]
        self.antenna_names = list(antenna_names[reordered_inds])
        antenna_positions = antenna_positions[reordered_inds, :]
        antenna_flags = antenna_flags[reordered_inds]
        cable_lens = cable_lens[reordered_inds]
        dig_gains = dig_gains[reordered_inds, :]

        # find flagged antenna
        flagged_ants = self.antenna_numbers[np.where(antenna_flags == 1)]

        # set parameters from other parameters
        self.Nants_data = len(self.antenna_numbers)
        self.Nants_telescope = len(self.antenna_numbers)
        self.Nbls = int(
            len(self.antenna_numbers) * (len(self.antenna_numbers) + 1) / 2.0
        )

        # get telescope parameters
        self.set_telescope_params()

        # build time array of centers
        time_array = np.arange(
            start_time + int_time / 2.0, end_time + int_time / 2.0 + int_time, int_time
        )

        # convert to time to jd floats
        float_time_array = Time(time_array, format="unix", scale="utc").jd.astype(float)
        # build into time array
        self.time_array = np.repeat(float_time_array, self.Nbls)

        self.Ntimes = len(time_array)

        self.Nblts = int(self.Nbls * self.Ntimes)

        # convert times to lst
        proc = self.set_lsts_from_time_array(background=background_lsts)

        self.integration_time = np.full((self.Nblts), int_time)

        # convert antenna positions from enu to ecef
        # antenna positions are "relative to
        # the centre of the array in local topocentric \"east\", \"north\",
        # \"height\". Units are meters."
        antenna_positions_ecef = uvutils.ECEF_from_ENU(
            antenna_positions, *self.telescope_location_lat_lon_alt
        )
        # make antenna positions relative to telescope location
        self.antenna_positions = antenna_positions_ecef - self.telescope_location

        # make initial antenna arrays, where ant_1 <= ant_2
        # itertools.combinations_with_replacement returns
        # all pairs in the range 0...Nants_telescope
        # including pairs with the same number (e.g. (0,0) auto-correlation).
        # this is a little faster than having nested for-loops moving over the
        # upper triangle of antenna-pair combinations matrix.
        ant_1_array, ant_2_array = np.transpose(
            list(
                itertools.combinations_with_replacement(
                    np.arange(self.Nants_telescope), 2
                )
            )
        )

        self.ant_1_array = np.tile(np.array(ant_1_array), self.Ntimes)
        self.ant_2_array = np.tile(np.array(ant_2_array), self.Ntimes)

        self.baseline_array = self.antnums_to_baseline(
            self.ant_1_array, self.ant_2_array
        )

        # create self.uvw_array
        self.set_uvws_from_antenna_positions(allow_phasing=False)

        # coarse channel mapping:
        # channels in group 0-128 go in order; channels in group 129-155 go in
        # reverse order
        # that is, if the lowest channel is 127, it will be assigned to the
        # first file
        # channel 128 will be assigned to the second file
        # then the highest channel will be assigned to the third file
        # and the next hightest channel assigned to the fourth file, and so on
        count = np.count_nonzero(coarse_chans <= 128)
        # map all file numbers to coarse channel numbers
        file_nums_to_coarse = {
            i + 1: coarse_chans[i]
            if i < count
            else coarse_chans[(len(coarse_chans) + count - i - 1)]
            for i in range(len(coarse_chans))
        }
        # map included coarse channels to file numbers
        coarse_to_incl_files = {}
        for i in included_file_nums:
            coarse_to_incl_files[file_nums_to_coarse[i]] = i
        # sort included coarse channels
        included_coarse_chans = sorted(coarse_to_incl_files.keys())
        # map included file numbers to an index that orders them
        file_nums_to_index = {}
        for i in included_coarse_chans:
            file_nums_to_index[coarse_to_incl_files[i]] = included_coarse_chans.index(i)
        # check that coarse channels are contiguous.
        chans = np.array(included_coarse_chans)
        for i in np.diff(chans):
            if i != 1:
                warnings.warn("coarse channels are not contiguous for this observation")
                break

        # warn user if not all coarse channels are included
        if len(included_coarse_chans) != len(coarse_chans):
            warnings.warn("some coarse channel files were not submitted")

        # build frequency array
        self.Nfreqs = len(included_coarse_chans) * num_fine_chans
        # TODO: Spw axis to be collapsed in future release
        self.freq_array = np.zeros((1, self.Nfreqs))

        # each coarse channel is split into 128 fine channels of width 10 kHz.
        # The first fine channel for each coarse channel is centered on the
        # lower bound frequency of that channel and its center frequency is
        # computed as fine_center = coarse_channel_number * 1280-640 (kHz).
        # If the fine channels have been averaged (added) by some factor, the
        # center of the resulting channel is found by averaging the centers of
        # the first and last fine channels it is made up of.
        # That is, avg_fine_center=(lowest_fine_center+highest_fine_center)/2
        # where highest_fine_center=lowest_fine_center+(avg_factor-1)*10 kHz
        # so avg_fine_center=(lowest_fine_center+lowest_fine_center+(avg_factor-1)*10)/2
        #                   =lowest_fine_center+((avg_factor-1)*10)/2
        #                   =lowest_fine_center+offset
        # Calculate offset=((avg_factor-1)*10)/2 to build the frequency array
        avg_factor = self.channel_width / 10000
        width = self.channel_width / 1000
        offset = (avg_factor - 1) * 10 / 2.0

        for i in range(len(included_coarse_chans)):
            # get the lowest fine freq of the coarse channel (kHz)
            lower_fine_freq = included_coarse_chans[i] * 1280 - 640
            # find the center of the lowest averaged channel
            first_center = lower_fine_freq + offset
            # add the channel centers for this coarse channel into
            # the frequency array (converting from kHz to Hz)
            self.freq_array[
                0, int(i * num_fine_chans) : int((i + 1) * num_fine_chans)
            ] = (
                np.arange(first_center, first_center + num_fine_chans * width, width)
                * 1000
            )
        # polarizations are ordered yy, yx, xy, xx
        self.polarization_array = np.array([-6, -8, -7, -5])

        if read_data:
            # read data into an array with dimensions (time, uv, baselines*pols)
            self.data_array = np.zeros(
                (self.Ntimes, self.Nfreqs, self.Nbls * self.Npols),
                dtype=data_array_dtype,
            )
            self.nsample_array = np.zeros(
                (self.Ntimes, self.Nbls, self.Nfreqs, self.Npols),
                dtype=nsample_array_dtype,
            )
            self.flag_array = np.full(
                (self.Ntimes, self.Nbls, len(included_coarse_chans), self.Npols), True
            )

            # read data files
            for file in file_dict["data"]:
                # get the file number from the file name
                file_num = int(file.split("_")[-2][-2:])
                # map file number to frequency index
                freq_ind = file_nums_to_index[file_num] * num_fine_chans
                with fits.open(
                    file, memmap=False, do_not_scale_image_data=False
                ) as hdu_list:
                    # count number of times
                    end_list = len(hdu_list)
                    for i in range(1, end_list):
                        time = (
                            hdu_list[i].header["TIME"]
                            + hdu_list[i].header["MILLITIM"] / 1000.0
                            + int_time / 2.0
                        )
                        time_ind = np.where(time_array == time)[0][0]
                        # dump data into matrix
                        # and take data from real to complex numbers
                        self.data_array[
                            time_ind, freq_ind : freq_ind + num_fine_chans, :
                        ] = (hdu_list[i].data[:, 0::2] + 1j * hdu_list[i].data[:, 1::2])
                        self.nsample_array[
                            time_ind, :, freq_ind : freq_ind + num_fine_chans, :
                        ] = 1.0
                        self.flag_array[
                            time_ind, :, file_nums_to_index[file_num], :
                        ] = False

            # build mapper from antenna numbers and polarizations to pfb inputs
            corr_ants_to_pfb_inputs = {}
            for i in range(len(antenna_numbers)):
                for p in range(2):
                    corr_ants_to_pfb_inputs[(antenna_numbers[i], p)] = 2 * i + p

            # for mapping, start with a pair of antennas/polarizations
            # this is the pair we want to find the data for
            # map the pair to the corresponding pfb input indices
            # map the pfb input indices to the pfb output indices
            # these are the indices for the data corresponding to the initial
            # antenna/pol pair
            # generate a mapping index array
            map_inds = np.zeros((self.Nbls * self.Npols), dtype=np.int32)
            # generate a conjugation array
            conj = np.full((self.Nbls * self.Npols), False, dtype=np.bool_)
            pfb_inputs_to_outputs = input_output_mapping()

            map_inds, conj = _corr_fits.generate_map(
                corr_ants_to_pfb_inputs, pfb_inputs_to_outputs, map_inds, conj,
            )

            # propagate coarse flags
            if propagate_coarse_flags:
                self.flag_array = np.any(self.flag_array, axis=2)
                self.flag_array = self.flag_array[:, :, np.newaxis, :]
                self.flag_array = np.repeat(self.flag_array, self.Nfreqs, axis=2)
            else:
                self.flag_array = np.repeat(self.flag_array, num_fine_chans, axis=2)

            # reorder data
            self.data_array = np.take(self.data_array, map_inds, axis=2)

            # conjugate data
            self.data_array[:, :, conj] = np.conj(self.data_array[:, :, conj])

            # reshape data
            self.data_array = self.data_array.reshape(
                (self.Ntimes, self.Nfreqs, self.Nbls, self.Npols)
            )

            self.data_array = np.swapaxes(self.data_array, 1, 2)

            # generage baseline flags for flagged ants
            bad_ant_inds = np.nonzero(
                np.logical_or(
                    np.in1d(ant_1_array, flagged_ants),
                    np.in1d(ant_2_array, flagged_ants),
                )
            )[0]
            self.flag_array[:, bad_ant_inds, :, :] = True

            # combine baseline and time axes
            self.data_array = self.data_array.reshape(
                (self.Nblts, self.Nfreqs, self.Npols)
            )
            self.flag_array = self.flag_array.reshape(
                (self.Nblts, self.Nfreqs, self.Npols)
            )
            self.nsample_array = self.nsample_array.reshape(
                (self.Nblts, self.Nfreqs, self.Npols)
            )

            # divide out digital gains
            if remove_dig_gains:
                # get gains for included coarse channels
                coarse_inds = np.where(np.isin(coarse_chans, included_coarse_chans))[0]
                # during commissioning a shift in the bit selection in the digital
                # receiver was implemented which effectively multiplies the data by
                # a factor of 64. To account for this, the digital gains are divided
                # by a factor of 64 here. For a more detailed explanation, see PR #908.
                dig_gains = dig_gains[:, coarse_inds] / 64
                dig_gains = np.repeat(dig_gains, num_fine_chans, axis=1)
                dig_gains1 = dig_gains[self.ant_1_array, :]
                dig_gains2 = dig_gains[self.ant_2_array, :]
                dig_gains1 = dig_gains1[:, :, np.newaxis]
                dig_gains2 = dig_gains2[:, :, np.newaxis]
                if self.data_array.dtype != np.complex128:
                    self.data_array = self.data_array.astype(np.complex128)
                self.data_array = self.data_array / (dig_gains1 * dig_gains2)
                if self.data_array.dtype != data_array_dtype:
                    self.data_array = self.data_array.astype(data_array_dtype)
            # divide out coarse band shape
            if remove_coarse_band:
                # get coarse band shape
                with open(
                    DATA_PATH + "/mwa_config_data/MWA_rev_cb_10khz_doubles.txt", "r"
                ) as f:
                    cb = f.read().splitlines()
                cb_array = np.array(cb).astype(np.float64)
                cb_array = cb_array.reshape(int(128 / avg_factor), int(avg_factor))
                cb_array = np.average(cb_array, axis=1)
                cb_array = cb_array[0:num_fine_chans]
                cb_array = np.tile(cb_array, len(included_coarse_chans))
                cb_array = cb_array[:, np.newaxis]
                if self.data_array.dtype != np.complex128:
                    self.data_array = self.data_array.astype(np.complex128)
                self.data_array = self.data_array / cb_array
                if self.data_array.dtype != data_array_dtype:
                    self.data_array = self.data_array.astype(data_array_dtype)

            # cable delay corrections
            if correct_cable_len:
                self.correct_cable_length(cable_lens)

            # add spectral window index
            self.data_array = self.data_array[:, np.newaxis, :, :]
            self.flag_array = self.flag_array[:, np.newaxis, :, :]
            self.nsample_array = self.nsample_array[:, np.newaxis, :, :]

            # because of an annoying discrepancy between file conventions, in order
            # to be consistent with the uvw vector direction, all the data must
            # be conjugated
            self.data_array = np.conj(self.data_array)

        # wait for LSTs if set in background
        if proc is not None:
            proc.join()

        if not self.metadata_only:
            # reorder polarizations
            # reorder pols calls check so must come after
            # lst thread is re-joined.
            self.reorder_pols()

        # phasing
        if phase_to_pointing_center:
            self.phase(ra_rad, dec_rad)

        if not self.metadata_only:
            if flag_init:
                self.flag_init(
                    num_fine_chans,
                    edge_width=edge_width,
                    start_flag=start_flag,
                    end_flag=end_flag,
                    flag_dc_offset=flag_dc_offset,
                )

            if use_cotter_flags:
                # throw an error if matching files not submitted
                if included_file_nums != included_flag_nums:
                    raise ValueError(
                        "flag file coarse bands do not match data file coarse bands"
                    )
                warnings.warn(
                    "coarse channel, start time, and end time flagging will default \
                        to the more aggressive of flag_init and AOFlagger"
                )
                for file in file_dict["flags"]:
                    flag_num = int(file.split("_")[-1][0:2])
                    # map file number to frequency index
                    freq_ind = file_nums_to_index[flag_num] * num_fine_chans
                    with fits.open(file) as aoflags:
                        flags = aoflags[1].data.field("FLAGS")
                    # some flag files are longer than data; crop the ends
                    flags = flags[: self.Nblts, :]
                    # some flag files are shorter than data; assume same end time
                    blt_ind = self.Nblts - len(flags)
                    flags = flags[:, np.newaxis, :, np.newaxis]
                    self.flag_array[
                        blt_ind:, :, freq_ind : freq_ind + num_fine_chans, :
                    ] = np.logical_or(
                        self.flag_array[
                            blt_ind:, :, freq_ind : freq_ind + num_fine_chans, :
                        ],
                        flags,
                    )

            # check if object is self-consistent
            # uvws are calcuated using pyuvdata, so turn off the check for speed.
            if run_check:
                self.check(
                    check_extra=check_extra,
                    run_check_acceptability=run_check_acceptability,
                    strict_uvw_antpos_check=strict_uvw_antpos_check,
                )


# -*- coding: utf-8 -*-
"""
CC相关的tags:
        某业务的业务模块
        获取idc和idc_unit(2个下拉)
ijobs相关的tags:
        某业务的ijobs作业列表, 如果有值，则选中(2个下拉);
"""
from operator import itemgetter

from blueking.component.ieod.shortcuts import get_client_by_request, get_client_by_user

from common.mymako import render_mako_tostring, render_json
from common.constant.server_attr import DEVICE_CLASS_CHOICES, LOGIC_ZONE_CHOICES, RAID_TYPE_CHOICES, OS_CHOICES

# ===============================================================================
# 在模板中以namespace方式引入, 则会自动传入模板原来的context作为第一个参数
# ===============================================================================
def server_attr(context, device_class=None, logic_zone=None, raid=None, os_name=None, server_num=0):
    '''
    服务器属性
    @usage: 在模板页面加入 ：
            <%namespace name="tags" module="test_app_l123.tags"/>
            然后添加：
                ${ tags.server_attr() }
            or
                ${ tags.server_attr(raid="No Raid") }
            or
                ${ tags.server_attr("C1", u"普通区", "No Raid", "SUSE Linux Enterprise Server 10 (x86_64)", 2) }
    '''
    # 下拉选项
    ctx = {'DEVICE_CLASS_CHOICES': DEVICE_CLASS_CHOICES,        # 设备类型
           'LOGIC_ZONE_CHOICES': LOGIC_ZONE_CHOICES,            # 逻辑区域
           'RAID_TYPE_CHOICES': RAID_TYPE_CHOICES,              # RAID结构
           'OS_CHOICES': OS_CHOICES,                            # 操作系统
           }
    # 如果有值，则进行赋值
    ctx.update({'device_class': device_class,
                'logic_zone': logic_zone,
                'raid': raid,
                'os_name': os_name,
                'server_num': server_num,
                })
    server_str = render_mako_tostring('/test_app_l123/tags/server_attr.html', ctx)
    return server_str

# 全局变量
# ===============================================================================
# app_id = 127              # 业务id
# ===============================================================================

def set_template_select(context, style_class, cc_app_id, is_template, template_id=None, level=None):
    """
    某业务的业务模块（下拉），在CC--Set实例化的时候需要用到"业务模块"这个参数
    @param style_class: 样式的class
    @param cc_app_id: cc系统的业务id
    @param is_template: 是否为模板页
    @param template_id: 选中的template id
    @param level: set模板的样式级别
    @usage: 在模板页面加入 ： <%namespace name="tags" module="test_app_l123.tags"/>
                                    然后添加：
                            ${ tags.set_template_select('', 127, 1) }           无初始值
                        or
                            ${ tags.set_template_select('span5', 127, 1, 1, 23) }  有初始值
    """
    username = context._kwargs['user']
    client = get_client_by_user(username)
    _template_list_result = client.cc.get_template_tree({'app_id': cc_app_id})
    _template_list = []
    if _template_list_result['result']:
        _template_list = _template_list_result['data']

    template_dict = {}
    template_list = []

    # 筛选数据
    if _template_list:
        for template in _template_list:
            template_list.append({'id': template['id'],
                                  'name': template['name'],
                                  })
            # 将template的id和 children做为键值对加入到字典中
            template_dict[template['id']] = template['children']
    # 返回参数
    ctx = {
        'template_dict': template_dict,       # set模板字典，key为template['id']，value为children（服务器规格信息）
        'template_list': template_list,       # set模版列表 [{'id': '33', 'name': u"作业名"},...]   注意，这里的id是字符串
        'template_id': template_id,           # 选中的模版id (string)
        'is_template': is_template,  # 是否为模板页
        'style_class': style_class,           # 样式class
        'level': level,                       # set模板的样式级别
    }
    server_str = render_mako_tostring('/test_app_l123/tags/set_template_select.html', ctx)
    return server_str

def idc_and_unit_select(context, style_class, is_template, idc_id=None,
                        idc_unit_id=None, idc_level=None, idc_unit_level=None):
    """
    获取idc和idc_unit(2个下拉), 如果有值，则选中
    @param style_class: 样式的class
    @param is_template: 是否为模板页
    @param idc_id:      选中的IDC的id
    @param idc_unit_id: 选中机房管理单元的id
    @param idc_level:  模板页中“IDC”是否选中，开区页中“IDC”是否显示
    @param idc_unit_level: 模板页中“机房管理单元”是否选中，开区页中“机房管理单元”是否显示
    @usage: 在模板页面加入 ： <%namespace name="tags" module="test_app_l123.tags"/>
                                    然后添加：
                            ${ tags.idc_and_unit_select('span5', 1, 1, 53, 355, 1, 1) }    有初始值
                        or
                            ${ tags.idc_and_unit_select('') }                  无初始值
    """
    idc_name = ''           # IDC名
    idc_unit_name = ''      # IDC_unit名
    # 获取组件client
    username = context._kwargs['user']
    client = get_client_by_user(username)
    # idc列表
    idc_list_result = client.cc.get_idc_list({})

    idc_list = []
    if idc_list_result['result']:
        idc_list = idc_list_result['data']

    #  获取选中idc的nname
    for idc in idc_list:
        if idc['IDCID'] == idc_id:
            idc_name = idc['IDCName']
            break
    # 如果“idc”已经选择，则要查询该idc的所有idc_unit列表
    idc_unit_list = []
    if idc_id:
        idc_unit_list_result = client.cc.get_idc_unit_list({'idc_id': idc_id})
        if idc_unit_list_result['result']:
            idc_unit_list = idc_unit_list_result['data']

        # 获取选中idc-unit的name
        for idc_unit in idc_unit_list:
            if idc_unit['IDCUnitID'] == idc_unit_id:
                idc_unit_name = idc_unit['IDCUnitName']
                break
    # 下拉选项
    ctx = {
        # idc列表  [{"IDCID":"1","IDCName":"什么什么机房"},...]
        'idc_list': sorted(idc_list, key=itemgetter('IDCName')),
        # idc_unit列表  [{{"IDCUnitID":"1","IDCUnitName":"什么什么机房4楼"},...]
        'idc_unit_list': sorted(idc_unit_list, key=itemgetter('IDCUnitName')),
        'idc_id': idc_id,                            # 选中的idc id
        'idc_unit_id': idc_unit_id,                  # 选中的idc_unit id
        'idc_name': idc_name,                        # IDC名
        'idc_unit_name': idc_unit_name,              # IDC_unit名
        'style_class': style_class,                  # 样式class
        'is_template': is_template,                   # 是否为模板页
        'idc_level': idc_level,                      # 模板页中“IDC”是否选中，开区页中“IDC”是否显示
        'idc_unit_level': idc_unit_level,            # 模板页中“机房管理单元”是否选中，开区页中“机房管理单元”是否显示
        'SITE_URL': context.get('SITE_URL'),         # SITE_URL
    }
    server_str = render_mako_tostring('/test_app_l123/tags/idc_and_unit_select.html', ctx)
    return server_str

def task_for_business_select(context, style_class, ijobs_app_id, is_template, ijobs_taskt_id=None,
                             ijobs_taski_id=None, taskt_level=None, taski_level=None):
    '''
    某业务的ijobs作业列表, 如果有值，则选中
    @param style_class: 样式的class
    @param ijobs_app_id: ijobs系统中的业务id
    @param is_template: 是否为模板页
    @param ijobs_taskt_id: 选中的作业id
    @param ijobs_taski_id: 选中的作业的执行态的id
    @param taskt_level: 模板页中“作业名”是否选中，开区页中“作业名”是否显示
    @param taski_level: 模板页中“执行状态”是否选中，开区页中“执行状态”是否显示
    @usage: 在模板页面加入 ： <%namespace name="tags" module="test_app_l123.tags"/>
                                    然后添加：
                            ${ tags.task_for_business_select('span5', 53, 1, 355, 35, 1) }    有初始值
                        or
                            ${ tags.task_for_business_select('', 53, 1) }                  无初始值
    '''
    # 获取组件client
    username = context._kwargs['user']
    client = get_client_by_user(username)
    # 作业列表
    _task_list_result = client.ijobs2.get_taskts({'app_id': ijobs_app_id})
    _task_list = []
    if _task_list_result['result']:
        _task_list = _task_list_result['data']

    task_list = []
    for task in _task_list:
        task_list.append({'id': str(task['id']),            # 把整型换成str型
                          'name': task['name'],
                          })
    # 如果“作业”已经选择，则要查询该作业的执行态列表
    taski_list = []
    if ijobs_taskt_id:
        _taski_list_result = client.ijobs2.get_taskis({'ijobs_taskt_id': ijobs_taskt_id})
        _taski_list = []
        if _taski_list_result['result']:
            _taski_list = _taski_list_result['data']

        for taski in _taski_list:
            taski_list.append({'id': str(taski['id']),      # 把整型换成str型
                               'name': taski['name'],
                               })
    # 下拉选项
    ctx = {
        'task_list': task_list,                      # 作业列表 [{'id': "33", 'name': u"作业名"},...]
        'taski_list': taski_list,                    # 作业的执行态列表
        'ijobs_taskt_id': ijobs_taskt_id,            # 选中的作业id
        'ijobs_taski_id': ijobs_taski_id,            # 选中的作业执行态id
        'style_class': style_class,                  # 样式class
        'is_template': is_template,                  # 是否为模板页
        'taskt_level': taskt_level,                  # 模板页中“作业名”是否选中，开区页中“作业名”是否显示
        'taski_level': taski_level,                   # 模板页中“执行状态”是否选中，开区页中“执行状态”是否显示
        'SITE_URL': context.get('SITE_URL'),         # SITE_URL
    }
    server_str = render_mako_tostring('/test_app_l123/tags/task_for_business_select.html', ctx)
    return server_str

# ===============================================================================
# ijobs web数据接口
# ===============================================================================
def ijobs_get_taskis_by_taskt(request, taskt_id):
    '''
    get taski list by taskt_id
    '''
    client = get_client_by_request(request)
    _taski_list_result = client.ijobs2.get_taskis({"ijobs_taskt_id": taskt_id})
    taskt_list = []
    if _taski_list_result['result']:
        taskt_list = _taski_list_result['data']
    result = {'result': True, 'message': taskt_list}
    return render_json(result)

# ===============================================================================
# cc web数据接口
# ===============================================================================
def cc_get_idc_unit(request):
    '''
    get idc_unit list by idc_name, if the idc_name is loss,then get all idc_unit
    @return:  [{"IDCUnitID":"1","IDCUnitName":"666dfgdfgsd"},...]
    '''
    idc_name = request.GET.get("idc_name", '')
    # 获取所有idc_unit列表
    client = get_client_by_request(request)
    _idc_unit_list_result = client.cc.get_idc_unit_list({})
    _idc_unit_list = []
    if _idc_unit_list_result['result']:
        _idc_unit_list = _idc_unit_list_result['data']

    idc_unit_list = []
    # 筛选idc下的idc_unit
    if idc_name:
        for idc_unit in _idc_unit_list:
            if idc_unit['IDCUnitName'].startswith(idc_name):
                idc_unit_list.append(idc_unit)
    # 返回
    result = {'result': True,
              'message': sorted(idc_unit_list, key=itemgetter('IDCUnitName')),
              }
    return render_json(result)


print(f'''The functions which this calculator can perform are
1.Addition
2.Subtraction
3.Multiplication
4.Division''')

n = input("Enter your selection: ")

if n in ('1' , '2' , '3' , '4' , '5' , '6' , '7'):
    a = int(input('Enter a Number: '))
    b = int(input('Enter a Number: '))

    if n == '1':
        c = a+b
        print(f"{a} + {b} = {c}")
    elif n == "2":
        c = a-b
        print(f"{a} - {b} = {c}")
    elif n == '3':
        c = a*b
        print(f"{a} x {b} = {c}")
    elif n == '4':
        c = a/b
        print(f"{a} / {b} = {c}")
else:
    print("Invalid Input")


import os
import typing as t
from abc import ABC, abstractmethod

from pony import orm
from vkwave.api import API
from vkwave.bots import Keyboard
from vkwave.client import AIOHTTPClient

from jacob.database.utils import admin, chats, groups, lists, students
from jacob.database.utils.storages import managers
from jacob.services import chats as chat_utils

JSONStr = str
api_session = API(tokens=os.getenv("VK_TOKEN"), clients=AIOHTTPClient())
api = api_session.get_context()

class Keyboards(ABC):
    """Базовая клавиатура-навигатор между экранами выбора студентов."""

    @abstractmethod
    def __init__(self, admin_id: int):
        """
        Создание новой клавиатуры.

        Args:
            admin_id: Идентификатор администратора
        """
        self.admin_id = admin_id
        self.return_to = ""

    @abstractmethod
    def menu(self) -> str:
        """Главное меню функции, реализуется в подклассах."""
        pass

    @abstractmethod
    def submenu(self, half: int) -> str:
        alphabet = students.get_unique_second_name_letters_in_a_group(
            admin.get_active_group(self.admin_id),
        )
        half_len = len(alphabet) // 2
        halfs = alphabet[:half_len], alphabet[half_len:]

        kb = Keyboard()

        for letter in halfs[half]:
            if len(kb.buttons[-1]) == 4:
                kb.add_row()
            kb.add_text_button(letter, payload={"button": "letter", "value": letter})
        if kb.buttons[-1]:
            kb.add_row()
        kb.add_text_button("◀ Назад", payload={"button": self.return_to})

        return kb.get_keyboard()

    @orm.db_session
    @abstractmethod
    def students(self, letter: str) -> str:
        data = students.get_list_of_students_by_letter(self.admin_id, letter)
        selected = managers.MentionStorageManager(
            self.admin_id,
        ).get_mentioned_students()
        half_index = self._find_half_index_of_letter(letter)
        kb = Keyboard()
        for item in data:
            if len(kb.buttons[-1]) == 2:
                kb.add_row()
            label = " "
            if item.id in selected:
                label = "✅ "
            kb.add_text_button(
                text=f"{label}{item.last_name} {item.first_name}",
                payload={
                    "button": "student",
                    "student_id": item.id,
                    "letter": letter,
                    "name": f"{item.last_name} {item.first_name}",
                },
            )
        if kb.buttons[-1]:
            kb.add_row()
        kb.add_text_button(
            text="◀ Назад",
            payload={"button": "half", "half": half_index},
        )

    def _get_halves_of_alphabet(self) -> t.Tuple[t.List[str], t.List[str]]:
        """
        Создает половины алфавита из списка студентов.

        Returns:
            t.Tuple[t.List[str]]: Половины алфавита
        """
        alphabet_: t.List[str] = students.get_unique_second_name_letters_in_a_group(
            admin.get_active_group(self.admin_id),
        )
        half_len = len(alphabet_) // 2

        return alphabet_[:half_len], alphabet_[half_len:]

    def _find_half_index_of_letter(self, letter: str) -> t.Optional[int]:
        """
        Определяет в какой половине алфавита находится буква и возвращает ее индекс.

        Args:
            letter: Буква для проверки

        Returns:
            int: Индекс половины
        """
        halves = self._get_halves_of_alphabet()

        saved_index = None

        for index, half in enumerate(halves):
            if letter in half:
                saved_index = index

        return saved_index

class StudentsNavigator(ABC):
    """Базовый конструктор клавиатур-навигаторов между экранами выбора студентов."""

    @abstractmethod
    def __init__(self, admin_id: int):
        self.admin_id = admin_id

    @abstractmethod
    def render(self) -> Keyboards:
        return Keyboards(self.admin_id)

def alphabet(admin_id: int) -> Keyboard:
    """
    Генерирует фрагмент клавиатуры с половинами алфавита фамилиий студентов.

    Args:
        admin_id: Идентификатор администратора

    Returns:
        Keyboard: Фрагмент клавиатуры
    """
    kb = Keyboard()
    alphabet = students.get_unique_second_name_letters_in_a_group(
        admin.get_active_group(admin_id),
    )
    half_len = len(alphabet) // 2
    f_alphabet, s_alphabet = alphabet[:half_len], alphabet[half_len:]
    for index, half in enumerate([f_alphabet, s_alphabet]):
        if half[0] == half[-1]:
            title = f"{half[0]}"
        else:
            title = f"{half[0]}..{half[-1]}"
        kb.add_text_button(title, payload={"button": "half", "half": index})

    return kb

async def list_of_chats(api_context, admin_id: int):
    """
    Генерирует фрагмент клавиатуры со списком подключенных чатов.

    Args:
        api_context: Объект API ВК.
        admin_id: идентификатор пользователя

    Returns:
        Keyboard: Фрагмент клавиатуры
    """
    kb = Keyboard()
    with orm.db_session:
        chat_objects = chats.get_list_of_chats_by_group(
            admin.get_active_group(admin_id),
        )
        for chat in chat_objects:
            chat_title = await chat_utils.get_chat_name(
                api_context,
                chat.vk_id,
            )
            kb.add_text_button(
                chat_title,
                payload={
                    "button": "chat",
                    "chat_id": chat.id,
                },
            )
    return kb

def prompt() -> Keyboard:
    """
    Генерирует клавиатуру с подтверждением действия.

    Returns:
        Keyboard: Объект клавиатуры
    """
    kb = Keyboard()
    kb.add_text_button(text="✅ Подтвердить", payload={"button": "confirm"})
    kb.add_text_button(text="🚫 Отменить", payload={"button": "deny"})
    return kb

def empty() -> JSONStr:
    kb = Keyboard()

    return kb.get_empty_keyboard()

def cancel() -> Keyboard:
    """
    Генерирует клавиатуру для отмены действия.

    Returns:
        Keyboard: клавиатура
    """
    kb = Keyboard()

    kb.add_text_button("🚫 Отмена", payload={"button": "cancel"})

    return kb

def confirm_with_chat_update():

    kb = prompt()

    if kb.buttons[-1]:
        kb.add_row()

    chat_emoji = "📡"
    kb.add_text_button(
        text=f"{chat_emoji} Переключить чат",
        payload={"button": "chat_config"},
    )

    return kb.get_keyboard()

def cancel_with_cleanup():
    kb = cancel()

    kb.add_row()
    kb.add_text_button("Очистить", payload={"button": "edit_cleanup"})

    return kb.get_keyboard()

def subgroups(group_id: int):
    kb = Keyboard()

    request = list(filter(bool, groups.get_subgroups_in_group(group_id)))

    for subgroup in request:
        if len(kb.buttons[-1]) == 2:
            kb.add_row()
        kb.add_text_button(
            subgroup,
            payload={
                "button": "subgroup",
                "subgroup": subgroup,
            },
        )

    if kb.buttons[-1]:
        kb.add_row()
    kb.add_text_button(text="◀ Назад", payload={"button": "presets"})

    return kb.get_keyboard()

def academic_statuses(group_id: int):
    kb = Keyboard()

    request = list(filter(bool, groups.get_academic_statuses_in_group(group_id)))

    for ac in request:
        if len(kb.buttons[-1]) == 2:
            kb.add_row()
        kb.add_text_button(
            ac.description,
            payload={
                "button": "ac_status",
                "status": ac.id,
            },
        )

    if kb.buttons[-1]:
        kb.add_row()

    kb.add_text_button(text="◀ Назад", payload={"button": "presets"})

    return kb.get_keyboard()

def custom_presets(group_id: int):
    kb = Keyboard()

    request = lists.get_lists_of_group(group_id)

    with orm.db_session:
        for preset in request:
            if len(kb.buttons[-1]) == 2:
                kb.add_row()
            kb.add_text_button(
                preset.name,
                payload={
                    "button": "preset",
                    "preset": preset.id,
                },
            )

    if kb.buttons[-1]:
        kb.add_row()

    kb.add_text_button(text="◀ Назад", payload={"button": "presets"})

    return kb.get_keyboard()


from qgreenland.tasks.common.fetch import FetchCmrGranule
from qgreenland.tasks.common.raster import WarpRaster
from qgreenland.util.luigi import LayerPipeline

class ArcticVegetation(LayerPipeline):
    def requires(self):
        fetch_data = FetchCmrGranule(
            dataset_cfg=self.cfg['dataset'],
            source_cfg=self.cfg['source'],
        )  # ->
        return WarpRaster(
            requires_task=fetch_data,
            layer_id=self.layer_id
        )


#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# main.py work with coreV2_1.py
# newer version integrate module panda to improve computation time and ease statistic calculation.

"""
__author__ = "Marc Oggier"
__license__ = "GPL"
__version__ = "1.0"
__maintainer__ = "Marc Oggier"
__contact__ = "Marc Oggier"
__email__ = "marc.oggier@gi.alaska.edu"
__status__ = "development"
__date__ = "2016/03/28"
"""

import os
import sys
import datetime as dt
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as colors
import matplotlib.lines as mlines
import numpy as np
import pickle

sys.path.append('../sea_ice')
import seaice.corePanda as sic
import seaice.mbsPanda as mbs
import seaice.cicePanda as cice
import seaice.icdtools

# =====================================================================================================================#
# USER VARIABLE INPUT
# =====================================================================================================================#

# data path
myhost = os.uname()[1]
if myhost == 'arran':
    data_path = '/mnt/data_lvm/'
elif myhost == 'islay':
    data_path = '/mnt/data_local'

# for Barrow'
location = 'BRW'
ics_filepath = os.path.join(data_path, 'seaice', 'core/BRW-FY_MBS/data.txt')  # ice core
mbs_data_dir = os.path.join(data_path, 'seaice', 'MBS/BRW/aoncadis')  # mass balance data (ice thickness, temperature profile)
weather_input_path = os.path.join(data_path, 'seaice','CICE/2014-0724Barrow1948-2013/JINforcing.txt-leap_00_12_ext14')  # weather data (T air)
cice_model_output_path = os.path.join(data_path, 'seaice', 'CICE/2014-0724Barrow1948-2013/modelout.txt')  #
freeze_up_path = os.path.join(data_path, 'seaice', 'core/BRW-FY_MBS/freezup_dates.txt')

# FDD/TDD classification
growth_start_day = dt.datetime(1900, 9, 1)  # Sept 1, is set arbitrary to the new winter day
freezing_temperature = -1.8  # C

# output data
output_dir = os.path.join(os.path.expanduser('~'), 'Desktop/output/pandas/', location)
ics_data_pkl_path = os.path.join(output_dir, 'ics_obs_data.pkl')
mbs_data_pkl_path = os.path.join(output_dir, 'mbs_obs_data.pkl')
ics_model_data_pkl_path = os.path.join(output_dir, 'ics_mod_data.pkl')
weather_input_pkl_path = os.path.join(output_dir, 'weather_data.pkl')
cice_model_output_pkl_path = os.path.join(output_dir, 'mod_output_data.pkl')
dd_data_pkl_path = os.path.join(output_dir, 'dd_data.pkl')
processed_data = os.path.join(output_dir, 'all_processed_data.pkl')

# toggle flag
display_figure = 'y'
save_figure = 'y'
save_figure_format = 'png'
display_comment = 'n'

# matplotlib config
font = {'size': 12} #'family': 'normal', 'weight': 'bold',
matplotlib.rc('font', **font)

# =====================================================================================================================#
# VARIABLES INITIALISATION
# =====================================================================================================================#
# figure
plt.close()
flag_figure_number = 0
nan_value = 'nan'

# output directory validity
if not os.path.isdir(output_dir):
    os.makedirs(output_dir)

if display_figure is 'n':
    import importlib
    import matplotlib
    matplotlib.use('Agg')
    importlib.reload(plt)

if save_figure is 'y':
    figure_dir = os.path.join(output_dir, 'figures')
    if not os.path.isdir(output_dir):
        os.makedirs(figure_dir)

# =====================================================================================================================#
# DATA IMPORTATION and PREPROCESSING
# =====================================================================================================================#
# ice core data
print('processing data importation ...')
# ice core data
if os.path.isfile(ics_data_pkl_path):
    ics_data = pickle.load(open(ics_data_pkl_path, 'rb'))
    print('ice core data loaded')
else:
    ics_data = sic.import_src(ics_filepath, variables='state variables')
    with open(ics_data_pkl_path, 'wb') as f:
        pickle.dump(ics_data, f)

for ii_key in ics_data.keys():
    ics_years = []
    ics_days_dict = {}
    for ii_key in sorted(ics_data.keys()):
        if ics_data[ii_key].coring_date not in ics_days_dict.values():
            if ics_data[ii_key].coring_date.month < 9 and ics_data[ii_key].coring_date.year - 1 not in ics_years:
                ics_years.append(ics_data[ii_key].coring_date.year - 1)
            elif ics_data[ii_key].coring_date.year not in ics_years:
                ics_years.append(ics_data[ii_key].coring_date.year)
        ics_days_dict[ii_key] = ics_data[ii_key].coring_date
    ics_years = sorted(ics_years)
print('ice core data imported')

# import model data
if os.path.isfile(cice_model_output_pkl_path):
    [model_data, cice_model] = pickle.load(open(cice_model_output_pkl_path, 'rb'))
    print('model output data loaded')
else:
    model_data = np.genfromtxt(cice_model_output_path)
    columns = ['year', 'month', 'day', 'ice thickness', 'snow thickness'] + ['T'+str('%02i' %ii) for ii in range(20)] + ['S'+str('%02i' %ii) for ii in range(20)]
    import pandas as pd
    cice_model = pd.DataFrame(model_data, columns=columns)
    cice_model['ice thickness'] = cice_model['ice thickness']/100
    cice_model['snow thickness'] = cice_model['snow thickness']/100
    cice_model['date'] = pd.to_datetime(cice_model.year*10000+cice_model.month*100+cice_model.day,format='%Y%m%d')
    with open(cice_model_output_pkl_path, 'wb') as f:
        pickle.dump([model_data, cice_model], f)
    print('model output data imported')

# extract ice core from model data at same date as observation
if os.path.isfile(ics_model_data_pkl_path):
    ics_model = pickle.load(open(ics_model_data_pkl_path, 'rb'))
    print('ics model output data loaded')
else:
    ics_model = {}
    ics_model_temp = []
    for ii_day in set(ics_days_dict.values()):
        if ii_day not in ics_model_temp:
            print(ii_day)
            run = 'run: ' + cice_model_output_path.split('/')[-2]
            ic_cice = cice.import_ice_core(model_data, ii_day, location, run)
            if ic_cice is not None:
                ics_model[ic_cice.core_name] = ic_cice
            ics_model_temp.append(ii_day)
    with open(ics_model_data_pkl_path, 'wb') as f:
        pickle.dump(ics_model, f)
    del ics_model_temp
    print('ics model output data imported')

# weather data
if os.path.isfile(weather_input_pkl_path):
    weather_data = pickle.load(open(weather_input_pkl_path, 'rb'))
    print('weather data loaded')
else:
    weather_data = np.genfromtxt(weather_input_path)
    with open(weather_input_pkl_path, 'wb') as f:
        pickle.dump(weather_data, f)
    print('weather data imported')

# freezup date
freezup_dates_data = np.genfromtxt(freeze_up_path, dtype=float)
freezup_dates = mbs.freezup_date_of_year(freezup_dates_data, ics_years)
print('freezup data imported')

# import mass balance data
if os.path.isfile(mbs_data_pkl_path):
    mbs_data = pickle.load(open(mbs_data_pkl_path, 'rb'))
    print('mbs data loaded')
else:
    mbs_data = {}

    for ii_year in range(2006, ics_years[-1] + 2):
        mbs_name = 'BRW_MBS_'+ str(ii_year)+'.csv'
        mbs_path = os.path.join(mbs_data_dir, mbs_name)
        mbs_data[ii_year] = mbs.read(mbs_path)
    with open(mbs_data_pkl_path, 'wb') as f:
        pickle.dump(mbs_data, f)
    print('mbs data imported')

# import temperature profile from mbs:
if os.path.isfile(ics_data_pkl_path) and mbs_data:
    ics_data = pickle.load(open(ics_data_pkl_path, 'rb'))
    ics_list_temp = sorted([ii for ii in ics_data.keys()])
    for ii_core in ics_list_temp:
        ic_data = ics_data[ii_core]
        ics_mbs_core_name = 'mbs-'+ic_data.coring_date.strftime('%Y%m%d')
        if 2005 < ic_data.coring_date.year:
            if ics_mbs_core_name not in ics_data.keys():
                print(ii_core)
                location_temp = 'Mass Balance Site, Barrow'
                if not isinstance(ic_data.ice_thickness, (float, int)):
                    ice_thickness = ic_data.ice_thickness
                    comment = "ice thickness from  " + ic_data.core_name
                elif not ic_data.profiles.empty:
                    if not np.isnan(ic_data.profiles.ice_core_length.dropna().mean()):
                        ice_thickness = ic_data.profiles.ice_core_length.dropna().mean()
                        comment = "ice thickness from ice core length of" + ic_data.core_name
                    elif not np.isnan(ic_data.profiles[ic_data.profiles.variable == 'salinity'].y_low.dropna().mean()):
                        ice_thickness = ic_data.profiles[ic_data.profiles.variable == 'salinity'].y_low.max()
                        comment = "ice thickness from salinity profile of ice core" + ic_data.core_name
                    elif not np.isnan(ic_data.profiles[ic_data.profiles.variable == 'temperature'].y_mid.dropna().mean()):
                        ice_thickness = ic_data.profiles[ic_data.profiles.variable == 'temperature'].y_mid.max()
                        comment = "ice thickness from temperature profile of ice core" + ic_data.core_name
                else:
                    ice_thickness = np.nan
                    comment = None
                ic_mbs = mbs.import_core(mbs_data, ic_data.coring_date, location=location_temp,
                                         ice_thickness=ice_thickness, comment=comment, overridehi = 'n')
                if ic_mbs is not None:
                    ic_mbs.add_corenames(ic_data.core_collection)
                    ics_data[ic_mbs.core_name]=ic_mbs
                    ics_data[ii_core].add_corenames(ic_mbs.core_name)
            # else:
            #     ics_data[ii_core].add_corenames(ic_mbs.core_name)
    del ics_list_temp
    with open(ics_data_pkl_path, 'wb') as f:
        pickle.dump(ics_data, f)
    print('temperature profile from mbs data imported')

# =====================================================================================================================#
# DEGREE DAY (DD) MODEL WITH OBSERVATION
# =====================================================================================================================#
DD = {}
l_FDD = 0

# compute DD for each years
for ii_year in freezup_dates.keys():
    freezup_day = dt.datetime.fromordinal(freezup_dates[ii_year])
    seasonal_cycle_end_date = dt.datetime(ii_year, growth_start_day.month, growth_start_day.day)
    DD_temp = seaice.icdtools.DegreeDayModel(weather_data, 7, freezup_day, seasonal_cycle_end_date, Tfreeze=freezing_temperature, Pday=dt.datetime(2000, 5, 15), T_data_unit='K')
    DD.update(DD_temp)
print('FDD and TDD computed')

# look for unique FDD and TDD, and add
FDD_ics = []
TDD_ics = []
for ii_day in set(ics_days_dict.values()):
    if DD[ii_day][1] < 0 and DD[ii_day][1] not in TDD_ics:
        TDD_ics.append(DD[ii_day][1])
    else:
        FDD_ics.append(DD[ii_day][0])
print('DD computed')

# =====================================================================================================================#
# PROCESSING
# =====================================================================================================================#
if not os.path.isfile(processed_data):
    ics_obs_stack = seaice.corePanda.DD_fillup(seaice.corePanda.stack_core(ics_data), DD, freezup_dates)
    ics_mod_stack = seaice.corePanda.DD_fillup(seaice.corePanda.stack_core(ics_model), DD, freezup_dates)
    print('stacking done')

    ics_obs_stack = ics_obs_stack.discretize(comment='y', display_figure='n')
    ics_mod_stack = ics_mod_stack.discretize(comment='y', display_figure='n')
    print('discretization done')

    # creating a 2D space depth (y) and time (TDD/FDD)
    bins_DD = np.unique(np.array([-100, -50, -35, -25, -15, -10, -5, 0, 250, 350, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000]))
    bins_y = np.linspace(0, 2, 41)
    stats = ['mean', 'std', 'min', 'max', 'median', 'mad']
    variables = ['salinity', 'temperature']
    for f_variable in variables:
        ics_obs_stack[f_variable] = ics_obs_stack[f_variable].astype(float)
        ics_mod_stack[f_variable] = ics_mod_stack[f_variable].astype(float)

    mod_grouped, mod_grouped_dict = ics_mod_stack.grouped_stat(variables, stats, bins_DD, bins_y)
    obs_grouped, obs_grouped_dict = ics_obs_stack.grouped_stat(variables, stats, bins_DD, bins_y)
    print('statistical grouped computed')

    with open(processed_data, 'wb') as f:
        pickle.dump([obs_grouped, obs_grouped_dict, ics_obs_stack, mod_grouped, mod_grouped_dict, ics_mod_stack, bins_DD, bins_y, variables, stats], f)

else:
    # import data from processed file data
    [obs_grouped, obs_grouped_dict, ics_obs_stack, mod_grouped, mod_grouped_dict, ics_mod_stack, bins_DD, bins_y, variables, stats] = pickle.load(open(processed_data, 'rb'))
    print('processed data imported')

# =====================================================================================================================#
# PLOTTING
# =====================================================================================================================#
# ## statistical envelop
DD_select = range(bins_DD.__len__()-1)
DD_label = [str(bins_DD[ii_bin])+' to '+str(bins_DD[ii_bin+1]) for ii_bin in DD_select]

for ii_bin in range(DD_select.__len__()):
    flag_figure_number += 1
    fig = plt.figure(flag_figure_number)
    fig.suptitle(DD_label[ii_bin])
    ax = []
    for ii_ax in range(variables.__len__()):
        f_variable = variables[ii_ax]
        ax.append(plt.subplot(1, len(variables), ii_ax + 1))
        mod_grouped.plot_stat_median(fig.axes[ii_ax], f_variable, DD_select[ii_bin])
        #obs_grouped.plot_stat_mean(fig.axes[ii_ax], f_variable, DD_select[ii_bin])
    plt.legend()

## Plot independent curve
DD_select = range(bins_DD.__len__()-1)
DD_label = [str(bins_DD[ii_bin])+' to '+str(bins_DD[ii_bin+1]) for ii_bin in DD_select]

import itertools

def flip(items, ncol):
    return itertools.chain(*[items[i::ncol] for i in range(ncol)])

for ii_bin in range(DD_select.__len__()):
    print(DD_label[ii_bin])
    flag_figure_number += 1
    fig = plt.figure(flag_figure_number, figsize=(12, 8), dpi=80, facecolor='w', edgecolor='w')
    fig.suptitle(DD_label[ii_bin])
    ax = []

    for ii_ax in range(variables.__len__()):
        f_variable = variables[ii_ax]
        if ii_ax == 0:
            ax.append(plt.subplot(1, variables.__len__(), ii_ax+1))
        else:
            ax.append(plt.subplot(1, variables.__len__(), ii_ax+1, sharey =ax[0]))

        # if obs_grouped[(obs_grouped.variable == f_variable) & (obs_grouped.DD_index == DD_select[ii_bin])].__len__() != 0:
        # obs_grouped.plot_stat_mean(fig.axes[ii_ax], f_variable, DD_select[ii_bin])
        mod_grouped.plot_stat_mean(fig.axes[ii_ax], f_variable, DD_select[ii_bin])
        # obs_grouped.plot_stat_median(fig.axes[ii_ax], f_variable, DD_select[ii_bin])

        core_list = sorted(set([item for ff_var in variables for item in obs_grouped_dict[ff_var][ii_bin]]))
        color_dict = {}
        # same color for core in between graphs
        for i_core in range(core_list.__len__()):
            color_dict[core_list[i_core]] = cm.jet(i_core/core_list.__len__())
        # same color for core belonging to a same coring event
        color_dict = {}
        handles, legends = ax[ii_ax].get_legend_handles_labels()
        for i_core in range(core_list.__len__()):
            if core_list[i_core] not in color_dict.keys():
                color = cm.jet(i_core / core_list.__len__())
                color_dict[core_list[i_core]] = color
                # legend and handle:
                handles.append(mlines.Line2D([], [], color=color))
                legends.append(core_list[i_core].split('-')[0] + '-' + core_list[i_core].split('-')[1][:8])
                for core in ics_obs_stack[ics_obs_stack.core_name == core_list[i_core]].core_collection.iloc[0]:
                    if core not in color_dict.keys():
                        color_dict[core] = color

        for core in core_list:
            ics_obs_stack.plot_core_profile(core, ax=fig.axes[ii_ax], variable=f_variable, param_dict={'color': color_dict[core], 'linewidth': 1})

    fig.set_size_inches(9, 8, forward=True)
    # plt.tight_layout()
    box = []
    v_shrink = 0.8
    for ii_ax in range(ax.__len__()):
        if ii_ax == 0:
            if float(DD_label[ii_bin].split(' to ')[0]) < 0:
                ax[ii_ax].set_xlim(0, 8)
            else:
                ax[ii_ax].set_xlim(-0, 16)
            ax[ii_ax].yaxis.tick_left()
            ax[ii_ax].spines['right'].set_visible(False)

        else:
            if float(DD_label[ii_bin].split(' to ')[0]) < 0:
                ax[ii_ax].set_xlim(-10, 0)
            else:
                ax[ii_ax].set_xlim(-20, 0)
            ax[ii_ax].spines['left'].set_visible(False)
            ax[ii_ax].yaxis.set_label_position('right')
            ax[ii_ax].yaxis.tick_right()

        ax[ii_ax].xaxis.tick_top()
        ax[ii_ax].xaxis.set_label_position('top')
        ax[ii_ax].spines['bottom'].set_visible(False)

        box.append(ax[ii_ax].get_position())
        ax[ii_ax].set_position([box[ii_ax].x0, box[ii_ax].y0 + box[ii_ax].height * (1-v_shrink),
                 box[ii_ax].width, box[ii_ax].height * v_shrink])
        if ii_ax == 0:
            plt.legend(flip(handles, 4), flip(legends, 4), loc='upper left', ncol=4, bbox_to_anchor=(-1.45, -0.05, 1.45 + 1.25, 0), mode="expand", fontsize=11)

    if save_figure is 'y':
        plt.savefig(os.path.join(figure_dir, (location+'-mod-'+DD_label[ii_bin]+'.'+save_figure_format)), format=save_figure_format)

cmap = {'temperature':plt.cm.Reds, 'salinity':plt.cm.Blues}
vmax = {'temperature':[0, 0],  'salinity':[8, 16]}
vmin = {'temperature':[-10, -20], 'salinity':[0, 0]}

f_variable = 'salinity'
y_interp = bins_y
DD_interp = np.sort(np.concatenate([bins_DD, bins_DD]))[1:-1]
z = {}
f_stat = 'median'
for f_variable in variables:
    z_temp = np.array([bins_DD.__len__(), y_interp.__len__()])
    for f_DD in range(bins_DD.__len__()-1):
        if obs_grouped[(obs_grouped.variable == f_variable) & (obs_grouped.stats == f_stat) & (obs_grouped.DD_index == f_DD)].DD_label.empty:
            data = [np.nan]
        else:
            data = obs_grouped[(obs_grouped.variable == f_variable) & (obs_grouped.stats == f_stat) & (obs_grouped.DD_index == f_DD)][f_variable].tolist()
            data.append(data[-1])
        ii_data = data.__len__()
        while ii_data < y_interp.__len__():
            data.append(np.nan)
            ii_data += 1
        if f_DD == 0:
            z_temp = data
            z_temp = np.vstack([z_temp, data])
        else:
            z_temp = np.vstack([z_temp, data])
            z_temp = np.vstack([z_temp, data])
    z[f_variable] = z_temp
    del z_temp

# look for season
DD_season = []
f_zero = 1
for ii_DD in range(DD_interp.__len__()):
    if f_zero and  DD_interp[ii_DD] <= 0:
        ii_season = 0
        if DD_season.__len__() == 0:
            DD_season.append([])
        DD_season[ii_season].append(ii_DD)
        if DD_interp[ii_DD] == 0:
            f_zero = 0
    else:
        ii_season = 1
        if DD_season.__len__() == 0:
            DD_season.append([])
        if DD_season.__len__() == 1:
            DD_season.append([])
        DD_season[ii_season].append(ii_DD)

for ii_season in range(DD_season.__len__()):
    if DD_interp[DD_season[ii_season]][0]<0:
        DD_interp[DD_season[ii_season]] = -DD_interp[DD_season[ii_season]]
        DD_season[ii_season] = DD_season[ii_season][::-1]

for f_variable in variables:
    for ii_season in range(DD_season.__len__()):
        fig = plt.figure()
        plt.pcolormesh(DD_interp[DD_season[ii_season]], y_interp, z[f_variable][DD_season[ii_season]].transpose(), vmin=vmin[f_variable][ii_season], vmax=vmax[f_variable][ii_season],
                       cmap=cmap[f_variable])
        plt.colorbar()
        plt.ylim(max(plt.ylim()), 0)
        plt.xlabel('time [DD]')
        if save_figure is 'y':
            plt.savefig(
                os.path.join(figure_dir, (location +'-'+ f_variable + str('%02i' % ii_season) + '.' + save_figure_format)),
                format=save_figure_format)

## obs-mod
cmap = {'temperature':plt.cm.bwr, 'salinity':plt.cm.bwr}
vmax = {'temperature':[5, 5],  'salinity':[4, 4]}
vmin = {'temperature':[-5, -5], 'salinity':[-4, -4]}
f_variable = 'salinity'
y_interp = bins_y
DD_interp = np.sort(np.concatenate([bins_DD, bins_DD]))[1:-1]
z = {}
f_stat = 'median'

z={}
for f_variable in variables:
    z_temp = np.array([bins_DD.__len__(), y_interp.__len__()])
    for f_DD in range(bins_DD.__len__()-1):
        if mod_grouped[(mod_grouped.variable == f_variable) & (mod_grouped.stats == f_stat) & (mod_grouped.DD_index == f_DD)].DD_label.empty:
            data_mod = [np.nan]
        else:
            data_mod = mod_grouped[(mod_grouped.variable == f_variable) & (mod_grouped.stats == f_stat) & (mod_grouped.DD_index == f_DD)][f_variable].tolist()
            data_mod.append(data[-1])
        ii_data = data_mod.__len__()
        while ii_data < y_interp.__len__():
            data_mod.append(np.nan)
            ii_data += 1

        if obs_grouped[(obs_grouped.variable == f_variable) & (obs_grouped.stats == f_stat) & (obs_grouped.DD_index == f_DD)].DD_label.empty:
            data_obs = [np.nan]
        else:
            data_obs = obs_grouped[(obs_grouped.variable == f_variable) & (obs_grouped.stats == f_stat) & (obs_grouped.DD_index == f_DD)][f_variable].tolist()
            data_obs.append(data[-1])
        ii_data = data_obs.__len__()
        while ii_data < y_interp.__len__():
            data_obs.append(np.nan)
            ii_data += 1

        data = np.array(data_obs)-np.array(data_mod)
        if f_DD == 0:
            z_temp = data
            z_temp = np.vstack([z_temp, data])
        else:
            z_temp = np.vstack([z_temp, data])
            z_temp = np.vstack([z_temp, data])
    z[f_variable] = np.nan_to_num(z_temp)
    del z_temp

# look for season
DD_season = []
f_zero = 1
for ii_DD in range(DD_interp.__len__()):
    if f_zero and  DD_interp[ii_DD] <= 0:
        ii_season = 0
        if DD_season.__len__() == 0:
            DD_season.append([])
        DD_season[ii_season].append(ii_DD)
        if DD_interp[ii_DD] == 0:
            f_zero = 0
    else:
        ii_season = 1
        if DD_season.__len__() == 0:
            DD_season.append([])
        if DD_season.__len__() == 1:
            DD_season.append([])
        DD_season[ii_season].append(ii_DD)

for ii_season in range(DD_season.__len__()):
    if DD_interp[DD_season[ii_season]][0]<0:
        DD_interp[DD_season[ii_season]] = -DD_interp[DD_season[ii_season]]
        DD_season[ii_season] = DD_season[ii_season][::-1]

for f_variable in variables:
    for ii_season in range(DD_season.__len__()):
        fig = plt.figure()
        plt.pcolormesh(DD_interp[DD_season[ii_season]], y_interp, z[f_variable][DD_season[ii_season]].transpose(), vmin=vmin[f_variable][ii_season], vmax=vmax[f_variable][ii_season],
                       cmap=cmap[f_variable])
        cbar = plt.colorbar()
        cbar.set_label(f_variable + ' '+seaice.properties.si_prop_unit[f_variable])
        plt.ylim(max(plt.ylim()), 0)
        plt.ylabel('ice thickness [m]')
        plt.xlabel('time [DD]')
        if save_figure is 'y':
            plt.savefig(
                os.path.join(figure_dir, (location +'-diff-'+ f_variable + str('%02i' % ii_season) + '.' + save_figure_format)),
                format=save_figure_format)

# 2009 thin sea ice in May 15

## mbs when is freezup

hi_freezup = 0.01
cice_freezup ={}
f_ice = 0
for f_day in cice_model[cice_model.year>1998].date:
    if not cice_model[(cice_model.date == f_day) & (cice_model['ice thickness']>hi_freezup)].empty:
        if f_ice == 0:
            f_ice = 1
            cice_freezup[int(cice_model[(cice_model.date == f_day)].year.tolist()[0])+1] = f_day.to_datetime()
    else:
        f_ice = 0
DD_freezup = {}
dt_freezup = {}
for f_year in range(2000, 2013):
    cice_day = cice_freezup[f_year]
    obs_day = dt.datetime.fromordinal(freezup_dates[f_year])
    if cice_day > obs_day:
        DD_freezup[f_year] = seaice.icdtools.DegreeDayModel(weather_data, 7, obs_day, cice_day, Tfreeze=freezing_temperature, T_data_unit='K')[cice_day][0]
        dt_freezup[f_year] = (cice_freezup[f_year]-dt.datetime.fromordinal(freezup_dates[f_year]))/dt.timedelta(minutes=60*24)
    else:
        DD_freezup[f_year] = -seaice.icdtools.DegreeDayModel(weather_data, 7, cice_day, obs_day, Tfreeze=freezing_temperature, T_data_unit='K')[obs_day][0]
        dt_freezup[f_year] = (cice_freezup[f_year]-dt.datetime.fromordinal(freezup_dates[f_year]))/dt.timedelta(minutes=60*24)

fig = plt.figure()
ax = []
ax.append(plt.subplot(1, 1, 1))
ax[0].bar(list(DD_freezup.keys()), DD_freezup.values(), align='center', color='r')
ax.append(ax[0].twinx())
ax[1].plot(list(dt_freezup.keys()), list(dt_freezup.values()), 'ko')
ax[0].axes.set_xlim(1999,2013)
ax[0].axes.set_xlabel('time [year]')
ax[0].yaxis.tick_right()
ax[0].yaxis.set_label_position('right')
ax[0].axes.set_ylabel('degree days difference between observation and model')
ax[1].yaxis.tick_left()
ax[1].yaxis.set_label_position('left')
ax[1].axes.set_ylabel('temporal difference between observation and model [days]')
plt.savefig(os.path.join(figure_dir, (location+'_DD_difference_between_obs_cice.' + save_figure_format)), format=save_figure_format)

#### maximal amount of core oin cor ecollection
temp = {}
for f_core in ics_obs_stack.core_name.unique():
    temp[f_core] = ics_obs_stack[ics_obs_stack.core_name == f_core].iloc[0]['core_collection'].__len__()

def flatten(xs):
    res = []
    def loop(ys):
        for i in ys:
            if isinstance(i, list):
                loop(i)
            elif i is not None:
                res.append(i)
    loop(xs)
    return res

core_col = sorted(['BRW_CS-20100319A', 'BRW_CS-20000229A', 'BRW_CS-20100513B',  'BRW_CS-20040215A', 'BRW_CS-20040607A', 'BRW_CS-20120513A', 'BRW_CS-20050530', 'BRW_CS-20100611A', ])
for f_core_col in core_col:
    core_list = ics_obs_stack[ics_obs_stack.core_name == f_core_col]['core_collection'].iloc[0]
    data = seaice.corePanda.CoreStack()
    variable_marker = {'temperature':0, 'salinity':0}
    for f_core in core_list:
        temp = ics_obs_stack[ics_obs_stack.core_name == f_core]
        data = data.add_profiles(temp)
        for f_variable in temp.variable.unique():
            variable_marker[f_variable] += 1
    for f_variable in variables:
        data[f_variable] = data[f_variable].astype(float)
    # # uppermost
    # data = data[data.y_mid < 0.225]
    # # middle
    # ic_bottom = data.ice_core_length.mean() - 0.225
    # data = data[data.y_mid < ic_bottom][data.y_mid > 0.225]
    # # bottomost
    ic_bottom = data.ice_core_length.mean()-0.225
    data = data[data.y_mid > ic_bottom]
    data_grouped, data_grouped_dict = data.grouped_stat(variables, stats, bins_DD, bins_y)
    for ii in range(data_grouped_dict[f_variable].__len__()):
        if data_grouped_dict[f_variable][ii]:
            DD_index = ii
    # # uppermost
    # z = list(set(flatten(data_grouped[data_grouped.DD_index == DD_index][data_grouped.y_mid < 0.225][
    #                          data_grouped.variable == 'salinity'][data_grouped.stats == 'std'][
    #                          'core_collection'].tolist()))).__len__()
    # a = np.nanmedian(data_grouped[(data_grouped.stats == 'std')][data_grouped.y_mid < 0.225][(data_grouped.variable == 'salinity')]['salinity'])
    # b = list(set(flatten(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid < 0.225][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    # c = np.nanmedian(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid < 0.225][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['salinity'])
    # d = list(set(flatten(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid < 0.225][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    # e = np.nanmedian(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid < 0.225][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['salinity'])

    # # center
    # z = list(set(flatten(data_grouped[data_grouped.DD_index == DD_index][data_grouped.y_mid < ic_bottom][data_grouped.y_mid > 0.225][data_grouped.variable == 'salinity'][data_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    # a = np.nanmedian(data_grouped[(data_grouped.stats == 'std')][data_grouped.y_mid < ic_bottom][data_grouped.y_mid > 0.225][(data_grouped.variable == 'salinity')]['salinity'])
    # b = list(set(flatten(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid < ic_bottom][obs_grouped.y_mid > 0.225][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    # c = np.nanmedian(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid < ic_bottom][obs_grouped.y_mid > 0.225][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['salinity'])
    # d = list(set(flatten(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid < ic_bottom][mod_grouped.y_mid > 0.225][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    # e = np.nanmedian(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid < ic_bottom][mod_grouped.y_mid > 0.225][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['salinity'])

    # bottom most
    z = list(set(flatten(data_grouped[data_grouped.DD_index == DD_index][data_grouped.y_mid > ic_bottom][
                             data_grouped.variable == 'salinity'][data_grouped.stats == 'std'][
                             'core_collection'].tolist()))).__len__()
    a = np.nanmedian(data_grouped[(data_grouped.stats == 'std')][(data_grouped.variable == 'salinity')]['salinity'])
    b = list(set(flatten(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid > ic_bottom][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    c = np.nanmedian(obs_grouped[obs_grouped.DD_index == DD_index][obs_grouped.y_mid > ic_bottom][obs_grouped.variable == 'salinity'][obs_grouped.stats == 'std']['salinity'])
    d = list(set(flatten(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid > ic_bottom][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['core_collection'].tolist()))).__len__()
    e = np.nanmedian(mod_grouped[mod_grouped.DD_index == DD_index][mod_grouped.y_mid > ic_bottom][mod_grouped.variable == 'salinity'][mod_grouped.stats == 'std']['salinity'])

    print(f_core_col, z, a, b, c, d, e)


import time
from functools import wraps

def logger_time(ignore=False):
	def real_decorator(func):
		@wraps(func)
		def inner(*args, **kwargs):
			if ignore:
				return func(*args, ** kwargs)
			print("-----")
			print('"' + func.__name__ + '" ' + "called:")
			ts = time.time()
			result = func(*args, **kwargs)
			te = time.time()
			print("time span:  %2.2f  secs" %(te-ts) )
			print("-----")
			return result
		return inner
	return real_decorator

@logger_time(ignore=False)
def slp_for_2():
	time.sleep(2)

if __name__ == "__main__":
	# for test
	slp_for_2()


"""
    Machine learning models such as Matrix Factorization and Restricted
Boltzmann Machines. (The latter is a work in progress). These are the kernels
of the suggest system, all other classes are built around those to ease their
use.

"""

#--internal imports
from testutils import TESTING

#--built-in imports
from warnings import warn
import random
import abc

#--external imports
from scipy.spatial.distance import cdist
import numpy as np

rng = np.random

#--code...
class Model(abc.ABC):
    """
        Abstract class defining the behavior of recommendation models. All
    suggest systems are forced to follow the rules set by this class. These
    methods' behaviors are therefore something you can rely on.

    """

    @abc.abstractmethod
    def update(self, *args, **kwargs):
        """
            Updates model's knowledge with given data and learning parameters.
        Learning paramteters depend on the model (MF or neighborhood) that you
        use but also on the mix-ins that you wrap over your model.

        return-type: none

        """

        raise NotImplementedError

    @abc.abstractmethod
    def suggest(self, user=None):
        """
            Practical use of the model, return recommendation for all users
        except if user (index) is given. In this case it returns only the given
        user's recommendation.

        return-type: matrix of shape (nb_users, nb_docs)

        """

        raise NotImplementedError

    @abc.abstractmethod
    def test(self, data, truth):
        """
            Tests the model's accuracy. The smaller the result, the better the
        model performs. (This is evaluated differently depending on the model
        used, do not rely on this to compare MF and Neighborhood).

        return-type: float

        """

        raise NotImplementedError

class MatrixFactorization(Model):
    """
        Matrix Factorization is a model-based technique for recommender systems
    which infers user and item profiles given user item interactions, and uses
    those profiles to recommend items to users. The model is much more
    expressive than a neighborhood model because user and item profiles are
    influenced by all direct and indirect relationships with both other users
    and other items. This model has the additional benefit of creating data
    which can have uses outside of recommendation. (Document or user clustering
    based on the profiles inferred, for example).

    """
    key = 1
    maxiter = 200

    def __init__(self, *, nb_users, nb_items, vector_size=20):
        self._users = np.random.normal(size=(nb_users, vector_size))
        self._items = np.random.normal(size=(nb_items, vector_size))

        self._u_biases = np.zeros(nb_users)
        self._i_biases = np.zeros(nb_items)
        self._bias = 0

    def update(self, *args, **kwargs):
        """
            This method invokes the SGD method which handles the training of the
        model. Simple as that.

        """
        self._SGD(*args, **kwargs)

    def _SGD(self, data, epochs=None, lmbda=0, eta=2):
        """
            Gradient descent.
        Since this model is linear, gradient descent is simple:
        The gradient for the user is the the loss function's gradient with
        respect to the model's result, multiplied by the item. Same goes for the
        item (however this time multiplied by the user).

            We train on all elements of the user_item_matrix which are non zero,
        treating elements equal to zero as missing data. The goal is to be able
        to fill that missing data with the knowledge acquired from the given
        data. One can recommend using the filled missing data (if the model
        thinks a user should see a document, the user's vector multiplied by the
        document's vector will give a big result relative to other
        multiplications).

        """
        global TESTING

        self._UIM = data
        self.eta = eta
        if epochs is None:
            iters = self.maxiter
        else:
            iters = epochs

        costs = []
        indices = list(zip(*data.nonzero()))
        for e in range(iters):
            random.shuffle(indices)
            for i, j in indices:
                self._grad_update(i, j, lmbda)

            costs.append(self.test(self._UIM))
            costs = costs[-20:]

            if len(costs) == 20 and epochs is None:
                if np.mean(np.asarray(costs[:-1]) - np.asarray(costs[1:])) < 0:
                    break

            if TESTING:
                print(costs[-1])

    def _grad_update(self, i, j, lmbda):
        """
            Calculates cost function's gradient with respect to the model's
        results. Cutting the SGD functionality in multiple methods helps in
        defining mix-ins for different weight update rules and other
        regularisation methods.

        """

        user, item = self._users[i], self._items[j]
        result = user.dot(item) + self._bias + self._u_biases[i] + self._i_biases[j]
        grad = (result - self._UIM[i, j])

        if np.isnan(grad):
            raise ArithmeticError

        self._update_weights(grad, i, j, lmbda, user, item)

    def _update_weights(self, grad, i, j, lmbda, user, item):
        """
            Like above. This method does the actual updating of the weights.

        """

        grad = self.eta * grad
        self._u_biases -= grad
        self._i_biases -= grad
        self._bias -= grad

        self._users[i] -= grad * item + lmbda * user**2
        self._items[j] -= grad * user + lmbda * item**2

    def suggest(self, users=None):
        """
            Takes a list of indices. Returns recommendation for all users
        pointed to by the indices. If users is None, returns suggestions for
        all users.

        """

        if users is None: return self._users.dot(self._items.T) + self._u_biases[:, np.newaxis] \
                         + self._i_biases + self._bias
        return self._users[users].dot(self._items.T) + self._u_biases[users, np.newaxis] \
               + self._i_biases + self._bias

    def test(self, truth):
        """
            Returns Mean Squared Error on given data.

        """

        indices = truth.nonzero()
        results = MatrixFactorization.suggest(self)
        return ((results[indices] - truth[indices])**2).mean()

class AdaGradMF(MatrixFactorization):
    """
        Automatically adapts the learning rate in a per-dimension way. Learning
    becomes more effective, more consistent and considerably faster.

    """
    key = 2

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._u_historical = np.zeros(self._users.shape)
        self._i_historical = np.zeros(self._items.shape)
        self._b_hist = 0

    def _update_weights(self, grad, i, j, lmbda, user, item):
        """
            A bunch of math stuff. Divides the gradient by the norm of the
        history of all gradients before this one.

        """

        u_grad, i_grad = item * grad, user * grad
        self._u_historical[i] += u_grad**2
        self._i_historical[j] += i_grad**2
        self._b_hist += grad**2

        u_sqrt = np.sqrt(self._u_historical[i])
        i_sqrt = np.sqrt(self._i_historical[j])
        u_grad = self.eta * u_grad / (1e-10 + u_sqrt)
        i_grad = self.eta * i_grad / (1e-10 + i_sqrt)

        self._users[i] -= u_grad + self.eta * (lmbda / (1e-10 + u_sqrt)) * user**2
        self._items[j] -= i_grad + self.eta * (lmbda / (1e-10 + i_sqrt)) * item**2

        grad = self.eta * grad / (1e-10 + np.sqrt(self._b_hist))
        self._u_biases[i] -= grad
        self._i_biases[j] -= grad
        self._bias -= grad

class Dropout(MatrixFactorization):
    """
        Adapted regularization method from neural nets. Gives surprisingly good
    results.

    """
    key = 3

    def update(self, p=0.7, *args, **kwargs):
        self.p = p
        super().update(*args, **kwargs)

    def _grad_update(self, i, j, lmbda):
        """
            Simply drops out certain elements of the user and item vectors (the
        same indices are used for both) hoping that it will prevent the features
        described by those elements from adapting to the presence of other
        features.

        """

        mask = rng.binomial(
            n=1, size=self._users.shape[1], p=self.p)
        user, item = self._users[i]*mask, self._items[j]*mask
        result = user.dot(item) + self._bias + self._u_biases[i] + self._i_biases[j]
        grad = (result - self._UIM[i, j])

        if np.isnan(grad):
            raise ArithmeticError

        self._update_weights(grad, i, j, lmbda, user, item)

    def suggest(self, *args, **kwargs):
        """
            Since we dropped out a certain number of elements each time we
        trained, we need to compensate for the fact that the not-dropped out
        elements adapted themselves to be fewer, and therefore needed to get
        bigger. We now use all the elements, so divide them.

        """

        return super().suggest(*args, **kwargs) * self.p


import sci_funcdef as fd
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

def main():

    fd.ratios(comm, rank, size)

if __name__ == "__main__":
    main()