from django.http import JsonResponse
from django.shortcuts import render, get_object_or_404
from accounts.models import UserProfile, Vendor

from buying.models import Address

from .models import PostPrice

def order_info(request, user_id):
    user_profile = get_object_or_404(UserProfile, user=request.user)
    adrs = user_profile.address_set.filter(validity=1).all()
    return render(request, 'buying/order-info.html', {"up": user_profile, 'adrs': adrs})

def get_post_price(request, user_id):
    vendor = get_object_or_404(Vendor, user=request.user)
    try:
        d = request.POST

        pp = vendor.postprice_set.filter(province=d['province'],
                            city=d['city'],
                            type=d['stype']).first()

        return JsonResponse({'stat': 1, 'price': pp.price})
    except:
        return JsonResponse({'stat': 0, 'price': 'این شهر قبلا ثبت نشده است.'})

def set_post_price(request, user_id):
    vendor = get_object_or_404(Vendor, user=request.user)
    try:
        d = request.POST

        pp = vendor.postprice_set.filter(province=d['province'],
                                         city=d['city'],
                                         type=d['stype'])

        if pp and pp.first():
            pp = pp.first()
            pp.price = d['price']
        else:
            pp = PostPrice(province=d['province'],
                           city=d['city'],
                           type=d['stype'],
                           vendor=vendor,
                           price=d['price']
                           )

        pp.save()
        return JsonResponse({'stat': 1})
    except:
        return JsonResponse({'stat': 0})

def add_address(request, user_id):
    # print(type(request.POST['city']))
    user_profile = get_object_or_404(UserProfile, user=request.user)
    try:
        d = request.POST
        address = Address(province=d['province'],
                          city=d['city'],
                          address=d['address'],
                          zipcode=d['zipcode'],
                          phone_number=d['phone_number'],
                          user_profile=user_profile)

        address.save()
        return JsonResponse({'stat': 1, 'add_id': address.id})
    except:
        return JsonResponse({'stat': 0})


import nltk
import A
import numpy as np
from nltk.align import Alignment
from nltk.align import AlignedSent
from collections import defaultdict
import math
class BerkeleyAligner():

    def __init__(self, align_sents, num_iter):
        self.MIN_PROB = 1e-10
        self.t, self.q = self.train(align_sents, num_iter)
    # TODO: Computes the alignments for align_sent, using this model's parameters. Return
    #       an AlignedSent object, with the sentence pair and the alignments computed.
    def align(self, align_sent):
        l = len(align_sent.mots)
        m = len(align_sent.words)
        best_alignment = []
        for j, trg_word in enumerate(align_sent.mots):
            best_prob = self.MIN_PROB
            best_alignment_point = None
            for i, src_word in enumerate(align_sent.words):
                align_prob = self.t['f'][(src_word, trg_word)] * self.q['f'][(j+1, i+1, l,m)] + self.t['e'][(trg_word, src_word)] * self.q['e'][(i+1, j+1, m, l)]

#               align_prob = self.t['f'][(src_word, trg_word)] * self.q['f'][(j+1, i+1, l,m)]

                if align_prob > best_prob:
                    best_prob = align_prob
                    best_alignment_point = i
            if best_alignment_point != None:
                best_alignment.append((best_alignment_point, j))
        new_sent = AlignedSent(align_sent.words, align_sent.mots, Alignment(best_alignment))
#
#        for j, trg_word in enumerate(align_sent.words):
#            best_prob = (self.t['f'][(trg_word, 'None')] * self.q['f'][(0, j+1, l, m)])
#            best_prob = max(best_prob, self.MIN_PROB)
#            best_prob = self.MIN_PROB
#            best_alignment_point = None
#            for i, src_word in enumerate(align_sent.mots):
#                align_prob = self.t['f'][(trg_word, src_word)] * self.q['f'][(i+1, j+1, l, m)] * self.t['e'][(src_word, trg_word)] * self.q['e'][(j+1, i+1, m, l)]
#                align_prob =  self.t['f'][(trg_word, src_word)] * self.q['f'][(i+1, j+1, l, m)]
#                if align_prob > best_prob:
#                    best_prob = align_prob
#                    best_alignment_point = i
#            if best_alignment_point != None:
#                best_alignment.append((j, best_alignment_point))
#        new_sent = AlignedSent(align_sent.words, align_sent.mots, Alignment(best_alignment))
#        align_sent.alignment = Alignment(best_alignment)
        return new_sent
    # TODO: Implement the EM algorithm. num_iters is the number of iterations. Returns the
    # translation and distortion parameters as a tuple.
    def train(self, aligned_sents, num_iters):
        t_f = defaultdict( lambda : 0.0)
        t_e = defaultdict( lambda : 0.0)
        q_e = defaultdict( lambda : 0.0)
        q_f = defaultdict( lambda : 0.0)
        n_e = defaultdict( lambda : 0.0)
        n_f = defaultdict( lambda : 0.0)

        for sentence in aligned_sents:
            len_mots = len(sentence.mots)
            len_words = len(sentence.words)
            for e in set(sentence.words):
                n_e[e] += len_mots
                n_e['None'] += len_mots
            for f in set(sentence.mots):
                n_f[f] += len_words
                n_f['None'] += len_words
        # t_f = IBMModel1(aligned_sents, 10).translation_table

        for sentence in aligned_sents:
            for e in sentence.words:
                for f in sentence.mots:
#                    t_e[(f,e)] = 1.0 / float(n_f[f])
#                    t_f[(e,f)] = 1.0 / float(n_e[e])
#                    t_e[(f,e)] = 1.0 / float(n_e[e])
#                    t_f[(e,f)] = 1.0 / n_f[f]
                    t_e[(f,e)] = 1.0 / n_e['None']
                    t_f[(e,f)] = 1.0 / n_f['None']
                    t_e[(f, 'None')] = 1.0 / n_e['None']
                    t_f[(e, 'None')] = 1.0 / n_f['None']
            m = len(sentence.mots)
            l = len(sentence.words)
            for i in range(m+1):
                for j in range(l+1):
                    if i == 0 and j == 0:
                        continue
                    q_e[(j,i,l,m)] = 1.0 / float(l + 1)
                    q_f[(i,j,m,l)] = 1.0 / float(m + 1)
        for s in range(num_iters):
            c_f = defaultdict( lambda : 0.0)
            c_e = defaultdict( lambda : 0.0)
            total_e = defaultdict( float)
            total_f = defaultdict(float)
            for sentence in aligned_sents:
                m = len(sentence.mots)
                l = len(sentence.words)
                e = sentence.words
                f = sentence.mots
                total_e = np.zeros(m+1, dtype = np.float)
                total_f = np.zeros(l+1, dtype=  np.float)
                ws = ['None'] + e
                ms = ['None'] + f
                for i in range(1, l+1):
                    total_f[i] = 0
                    for j in range(0, m+1):
                        total_f[i] += q_f[(j, i, m, l )] * t_f[(ws[i], ms[j])]
                for j in range(1, m+1):
                    total_e[j] = 0
                    for i in range(0, l+1):
                        total_e[j] += q_e[(i, j, l, m)] * t_e[(ms[j], ws[i])]

#                for i in range( m + 1):
#                    for j in range(l + 1):
#                        if i != 0:
#                            total_e[i] += q_e[(j,i,l,m)] * t_e[(ms[i], ws[j])]
#                        if j != 0:
#                            total_f[j] += q_f[(i,j,m,l)] * t_f[(ws[j], ms[i])]
                        #total_mots[i] += q[(j,i,l,m)] * t[(f[i], e[j])]
                        #total_words[i] += q[(i,j,m,l)] * t[(e[j], f[i])]
                for i in range(m+1):
                    for j in range(l+1):
                        if i==0 and j == 0:
                            continue
#                        if i != 0:
                        if not total_e[i] == 0:
                            theta1 = q_e[(j,i,l,m)] * t_e[(ms[i],ws[j])] * 1.0 / total_e[i]
                        else:
                            theta1 = self.MIN_PROB
                        if not total_f[j] == 0:
                            theta2 = q_f[(i,j,m,l)] * t_f[(ws[j],ms[i])] * 1.0 / total_f[j]
                        else:
                            theta2 = self.MIN_PROB
                        theta = 0.5 * theta1 + 0.5 * theta2
#                        theta = math.sqrt(theta1 + theta2)
                        c_f[(ws[j], ms[i])] += theta
                        c_e[(ms[i], ws[j])] += theta
                        c_f[ms[i]] += theta
                        c_e[ws[j]] += theta
                        c_f[(i,j,m,l)] += theta
                        c_e[(j,i,l,m)] += theta
                        c_f[(j,m,l)] += theta
                        c_e[(i,l,m)] += theta

            ws = []
            ms = []
            for sentence in aligned_sents:
                ws = ['None'] + sentence.words
                ms = ['None'] + sentence.mots
                l = len(sentence.words)
                m = len(sentence.mots)

                for i in range(m+1):
                    for j in range(l+1):
                        if i == 0 and j == 0:
                            continue
                        if c_f[(i,j,m,l)] >0 and not c_f[(j,m,l)] == 0:
                            q_f[(i,j,m,l)] = 1.0 * c_f[(i,j,m,l)] / c_f[(j,m,l)]
                        if c_e[(j,i,l,m)] > 0 and not c_e[(i,l,m)] ==0:
                            q_e[(j,i,l,m)] = 1.0 * c_e[(j,i,l,m)] / c_e[(i,l,m)]

#                        if c_e[(i,l,m)] == 0:
#                            q_e[(j,i,l,m)] = c_e[(j,i,l,m)] * 1.0 / c_e[(i,l,m)]
#                        if not c_f[(j,m,l)] == 0:
#                            q_f[(i,j,m,l)] = c_f[(i,j,m,l)] * 1.0 / c_f[(j,m,l)]
                        if c_e[(ms[i], ws[j])] >0 and not c_e[ws[j]] == 0:
                            t_e[(ms[i], ws[j])] = 1.0 * c_e[(ms[i], ws[j])] / c_e[ws[j]]
                        if c_f[(ws[j], ms[i])] >0 and not c_f[ms[i]] == 0:
                            t_f[(ws[j], ms[i])] = 1.0 * c_f[(ws[j], ms[i])] / c_f[ms[i]]

#                        t_e[(ms[i], ws[j])] = 1.0 * c_e[(ms[i], ws[j])] / c_e[ws[j]]
#                        t_f[(ws[j], ms[i])] = 1.0 *  c_f[(ws[j], ms[i])] / c_f[ms[i]]
        t = {}
        q = {}
        t['f'] = t_f
        q['f'] = q_f
        t['e'] = t_e
        q['e'] = q_e
#        t = t_f
#        q =  q_f
        return (t,q)

def main(aligned_sents):
    ba = BerkeleyAligner(aligned_sents, 10)
    A.save_model_output(aligned_sents, ba, "ba.txt")
    avg_aer = A.compute_avg_aer(aligned_sents, ba, 50)

    print ('Berkeley Aligner')
    print ('---------------------------')
    print('Average AER: {0:.3f}\n'.format(avg_aer))


# type: ignore
"""Train model with parameters from config"""

import argparse
import os
import sys

from argus import load_model
import torch
import torch.distributed as dist
from timm.utils.model_ema import ModelEmaV2 as ModelEma

from modules.factory import get_loaders, get_model, get_callbacks, get_metrics, get_nn_params

sys.path.append("/workdir/src/")

def run(args):
    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        dist.init_process_group(backend='nccl', init_method='env://')

    train_loader, valid_loader = get_loaders(args.distributed, args.local_rank)
    callbacks, val_callbacks = get_callbacks(args.distributed, args.local_rank)
    metrics = get_metrics()
    if args.model is not None:
        model_params = get_nn_params()
        model = load_model(
            file_path=args.model,
            optimizer=model_params['optimizer'],
            loss=model_params['loss'],
            device=model_params['device']
        )
    else:
        model = get_model(args.distributed, args.local_rank)

    if args.ema:
        model.ema_model = ModelEma(
            model=model.get_nn_module(),
            decay=0.995,
            device=model_params['device']
        )

    model.fit(
        train_loader=train_loader,
        val_loader=valid_loader,
        num_epochs=args.num_epochs,
        metrics=metrics,
        metrics_on_train=True,
        callbacks=callbacks,
        val_callbacks=val_callbacks,
    )

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--num_epochs',
                        type=int,
                        default=500)

    parser.add_argument('--model', '-m',
                        type=str,
                        default=None,
                        help="path to saved model to continue from checkpoint")

    parser.add_argument('-ema',
                        action='store_true')

    parser.add_argument("--local_rank", default=0, type=int)

    arguments = parser.parse_args()
    arguments.distributed = False
    if 'WORLD_SIZE' in os.environ:
        arguments.distributed = int(os.environ['WORLD_SIZE']) > 1

    run(arguments)


import numpy as np
from scipy import stats
import datetime as dt
import pytz
import copy
import pandas as pd

from gravtools.models.misc import time_it

# optional imports:
try:
    import geopandas
except ImportError:
    _has_geopandas = False
else:
    _has_geopandas = True

from gravtools import settings

class LSM:
    """Base class for LSM classes.

    Attributes
    ----------
    lsm_method : str
        Defines the adjustment method. Has to b listed in :py:obj:`gravtools.settings.ADJUSTMENT_METHODS`.
    stat_df : :py:obj:`gravtools.Station.stat_df`
        The station dataframe contains all relevant station data.
    setups : dict of dicts
        The setups dictionary contains all observation data used for the adjustment. The keys of the dictionary
        are the survey names (str). The items are again keys with the following items:

        - ref_epoch_delta_t_h : datetime object
            Reference epoch for the relative reference times in the column `delta_t_h` in the `setup_df` dataframe.
            The reference epoch is determined as the epoch of the first (active) observation in this survey.
        -  ref_epoch_delta_t_campaign_h : datetime object
            Reference epoch for the relative reference times in the column `delta_t_campaign_h` in the `setup_df`
            dataframe. The reference epoch is determined as the epoch of the first (active) observation in the campaign.
        - setup_df : Pandas DataFrame
            Pandas dataframes containing the observation data (see :py:obj:`gravtools.Survey.setup_df`).
        - tide_correction_type : str, optional (default='')
            Type of the tidal corrections applied on the reduced observations (column `g_red_mugal` in `obs_df`) that
            were used to calculate the setup data store in this dict. Valid entries have to be listed in
            :py:obj:`gravtools.settings.TIDE_CORRECTION_TYPES`.
        - reference_height_type : str, optional (default='')
            Reference height type of the reduced observations (column `g_red_mugal` in `obs_df`) that
            were used to calculate the setup data store in this dict. Valid entries have to be listed in
            :py:obj:`gravtools.settings.REFERENCE_HEIGHT_TYPE`.

    comment : str, optional (default = '')
        Optional comment on the adjustment run.
    init_time : datetime object
        Representing the date and time the LSM object has been initialized.
    write_log : bool, optional (default=True)
        Flag that indicates whether log string should be written or not.
    log_str : str
        String to log the status of the adjustment.
    setup_obs_df : Pandas DataFrame
        Pandas Dataframes for logging (differential or absolute) setup observations, the related metadata, estimation
        results and statistics. The columns of the dataframe may differ between adjustment methods.
    number_of_iterations : int (default=0)
        Indicates the number of iterations if iterative adjustment was applied. `0` indicated that a non-iterative
        adjustment was applied.
    drift_ref_epoch_type : string ('survey' or 'campaign'), optional (default='survey')
        Defines whether the reference epoch t0 for the estimation of the drift polynomials for each survey in the
        campaign is the reference epoch of the first (active) observation in each survey (option: 'survey') or the
        first (active) observation in the whole campaign (option: 'campaign').
        """

    def __init__(self, lsm_method, stat_df, setups, comment='', write_log=True):
        """
        Parameters
        ----------
        lsm_method : str
            Defines the adjustment method. Has to b listed in :py:obj:`gravtools.settings.ADJUSTMENT_METHODS`.
        stat_df : :py:obj:`gravtools.Station.stat_df`
            The station dataframe contains all relevant station data.
        setups : dict of dicts
            The setups dictionary contains all observation data used for the adjustment. The keys of the dictionary
            are the survey names (str). The items are again keys with the following items:

            - ref_epoch_delta_t_h : datetime object
                Reference epoch for the relative reference times in the column `delta_t_h` in the `setup_df` dataframe.
                The reference epoch is determined as the epoch of the first (active) observation in this survey.
            -  ref_epoch_delta_t_campaign_h : datetime object
                Reference epoch for the relative reference times in the column `delta_t_campaign_h` in the `setup_df`
                dataframe. The reference epoch is determined as the epoch of the first (active) observation in the campaign.
            - setup_df : Pandas DataFrame
                Pandas dataframes containing the observation data (see :py:obj:`gravtools.Survey.setup_df`).
            - tide_correction_type : str, optional (default='')
                Type of the tidal corrections applied on the reduced observations (column `g_red_mugal` in `obs_df`) that
                were used to calculate the setup data store in this dict. Valid entries have to be listed in
                :py:obj:`gravtools.settings.TIDE_CORRECTION_TYPES`.
        - red_reference_height_type : str, optional (default='')
                Reference height type of the reduced observations (column `g_red_mugal` in `obs_df`) that
                were used to calculate the setup data store in this dict. Valid entries have to be listed in
                :py:obj:`gravtools.settings.REFERENCE_HEIGHT_TYPE`.

        comment : str, optional (default = '')
            Optional comment on the adjustment run.
        write_log : bool, optional (default=True)
            Flag that indicates whether log string should be written or not.

        Notes
        -----
        In order to prevent altering the original input data (in case of assignment via reference), deep copies of the
        input data are stored in objects of this class.
        """
        # Initial input checks:
        if lsm_method in settings.ADJUSTMENT_METHODS.keys():
            self.lsm_method = lsm_method
        else:
            raise ValueError(f'"{lsm_method}" is not a valid adjustment method identifier! All valid methods are '
                             f'defined in gravtools.settings.ADJUSTMENT_METHODS.')
        if isinstance(comment, str):
            self.comment = comment
        else:
            raise ValueError(f'"{comment}" need to be a string.')
        if isinstance(write_log, bool):
            self.write_log = write_log
        else:
            raise ValueError(f'"{write_log}" needs to be a boolean.')
        self.init_time = dt.datetime.now(tz=pytz.utc)
        self.log_str = ''

        # Create deep copies of the input data items:
        self.stat_df = stat_df.copy(deep=True)
        self.setups = copy.deepcopy(setups)

        # Initialize attributes:
        self.setup_obs_df = None  # Dataframe that contain the observation (setup) related results
        self.observed_stations = None  # Unique list of observed stations; defines the station IDs for matrices
        self.stat_obs_df = None  # Station dataframe that contains estimation results of observed stations
        self.drift_pol_df = None  # DataFrame that contains the estimated parameters of the drift polynomials and their statistics for each survey
        self.vg_pol_df = None  # DataFrame that contains the estimated parameters of the vertical gravity gradient polynomial

        # Estimation settings:
        self.drift_polynomial_degree = None
        self.sig0_a_priori = None  # A priori standard deviation of unit weight
        self.scaling_factor_datum_observations = None
        self.confidence_level_chi_test = None
        self.confidence_level_tau_test = None
        self.drift_ref_epoch_type = ''  # 'survey' or 'campaign'
        self.vg_polynomial_degree = None
        self.vg_polynomial_ref_height_offset_m = 0.0

        # General statistics:
        self.number_of_stations = None
        self.number_of_datum_stations = None
        self.number_of_estimates = None
        self.degree_of_freedom = None

        # A posteriori statistics:
        self.s02_a_posteriori = None  # A posteriori variance of unit weight

        # Matrices:
        self.Cxx = None  # Co-variance matrix of estimated parameters
        self.Rxx = None  # Correlation matrix of estimates parameters
        self.x_estimate_names = []  # List of items in the Cxx matrix (index 0 to n)
        self.mat_A = None  # Design matrix
        self.mat_x = None  # Estimated parameters

        # Iterative adjustment:
        self.number_of_iterations = 0  # `0` indicates no iterations.

        # Statistical tests:
        self.number_of_outliers = None
        self.global_model_test_status = ''  # str

    @classmethod
    def from_campaign(cls, campaign, comment='', write_log=True):
        """Constructor that generates and populates the LSM object (child!) from a Campaign class object.

        Notes
        -----
        Put all checks and data preparations relevant for all LSM methods into this method!

        Parameters
        ----------
        campaign : :py:obj:`gravtools.models.survey.Campaign`
            The campaign object needs to provide setup data for all active surveys and the related station data.
        comment : str, optional (default = '')
            Arbitrary comment on the LSM run.
        write_log : bool, optional (default=True)
            Flag that indicates whether log string should be written or not.

        Returns
        -------
        LSM Object of the child class
            Contains all information required for adjusting the campaign.
        """
        # Check if all required data is available in the campaign object:

        # Comment:
        if not isinstance(comment, str):
            raise TypeError(f'"comment" needs to be a string!')

        # Station data:
        if campaign.stations is None:
            raise AssertionError(f'The campaign "{campaign.campaign_name}" does not contain any station data!')
        else:
            if not hasattr(campaign.stations, 'stat_df'):
                raise AssertionError(f'The campaign "{campaign.campaign_name}" has not station dataframe!')
            else:
                if len(campaign.stations.stat_df) == 0:
                    raise AssertionError(f'The campaign "{campaign.campaign_name}" has an empty station dataframe!')

        # Survey data:
        if campaign.surveys is None:
            raise AssertionError(f'The campaign "{campaign.campaign_name}" does not contain any survey data!')
        else:
            if len(campaign.surveys) == 0:
                raise AssertionError(f'The campaign "{campaign.campaign_name}" contains no survey data!')

        # Create setups dict:
        # Loop over surveys in campaign:
        setups = {}
        for survey_name, survey in campaign.surveys.items():
            if survey.keep_survey and survey.is_active:
                if survey.setup_df is None:
                    raise AssertionError(f'Setup data is missing for survey "{survey_name}"')
                else:
                    if len(survey.setup_df) < 2:
                        raise AssertionError(f'Survey "{survey_name}" has less than two setup observations! '
                                             f'A minimum of two setups is required in order to define '
                                             f'differential observations!')
                    else:
                        setup_data_dict = {'ref_epoch_delta_t_h': survey.ref_delta_t_dt,
                                           'ref_epoch_delta_t_campaign_h': campaign.ref_delta_t_dt,
                                           'setup_df': survey.setup_df,
                                           'tide_correction_type': survey.setup_tide_correction_type,
                                           'reference_height_type': survey.setup_reference_height_type,
                                           'atm_pres_correction_type': survey.setup_atm_pres_correction_type,
                                           'setup_calc_method': survey.setup_calc_method,
                                           'setup_obs_list_df': survey.setup_obs_list_df}
                        setups[survey_name] = setup_data_dict

        # Check if setup data is available:
        if len(setups) == 0:
            raise AssertionError(f'Setup data of campaign "{campaign.campaign_name}" does not contain observations!')

        # Initialize and return LSM object:
        return cls(campaign.stations.stat_df, setups, comment=comment, write_log=write_log)

    def adjust_autoscale_s0(self,
                            iteration_approach='Multiplicative',
                            s02_target=1,
                            s02_target_delta=0.1,
                            max_number_iterations=10,
                            add_const_to_sd_of_observations_step_size_mugal=5.0,
                            max_total_additive_const_to_sd_mugal=20.0,
                            multiplicative_factor_step_size_percent=10.0,  # [%]
                            max_multiplicative_factor_to_sd_percent=200.0,  # [%]
                            min_multiplicative_factor_to_sd_percent=50.0,  # [%]
                            drift_pol_degree=1,
                            sig0_mugal=1,
                            scaling_factor_datum_observations=1.0,
                            add_const_to_sd_of_observations_mugal=0.0,
                            scaling_factor_for_sd_of_observations=1.0,
                            confidence_level_chi_test=0.95,
                            confidence_level_tau_test=0.95,
                            drift_ref_epoch_type='survey',
                            noise_floor_mugal=0.0,
                            verbose=False,
                            ):  # Confidence level):
        """Run the adjustment iteratively in order to adjust s0 to the target value by adapting the SD of observations.

        Notes
        -----
        Iterative adjustment works for differential and non-differential LSM adjustment, not for VG estimation!

        Parameters
        ---------
        iteration_approach : str, optional (default='Multiplicative')
            Defines the iteration approach (Multiplicative od Additive)
        s02_target : float, optional (default=1.0)
            Target a posteriori s0² (variance of unit weight) for the iteration.
        s02_target_delta : float, optional (default=0.1)
            Permissible deviation of the target a posteriori s0². As soon as the a posteriori s0 of the current lsm run
            lies withing the threshold of `s02_target` +- `s02_target_delta` the iteration was successful and stops.
        max_number_iterations : int, optional (default=10)
            Maximum allowed number of iterations. If the a posteriori s0 does not lie within the defined threshold
            (`s02_target` +- `s02_target_delta`) after `max_number_iterations` iterations an assertion error is raised.
        add_const_to_sd_of_observations_step_size_mugal : float, optional (default=5.0)
            Initial iteration step size for the additive constant that is added to the standard deviation of all
            observations. This parameter is only considered when using the `additive` iteration approach.
        max_total_additive_const_to_sd_mugal : float, optional (default=20.0)
            If the additive factor for the SD of observations that is determined iteratively in order to reach the
            defined target s0 is larger than `max_total_additive_const_to_sd_mugal`, the iteration procedure
            failed and an assertion error is raised. This parameter is only considered when using the `additive`
            iteration approach.
        multiplicative_factor_step_size_percent : float, optional (default=10.0)
            Initial iteration step size for the multiplicative factor that is used to scale all setup
            observations. This parameter is only considered when using the `multiplicative` iteration approach.
        max_multiplicative_factor_to_sd_percent : float, optional (default=200.0)
            Maximum scaling factor when using the `multiplicative` iteration approach for scaling the SD of setup
            observations. Minimum = 100%.
        min_multiplicative_factor_to_sd_percent : float, optional (default=50.0)
            Minimum scaling factor when using the `multiplicative` iteration approach for scaling the SD of setup
            observations. Minimum = 1.0%, maximum = 100.0%
        drift_pol_degree : int, optional (default=1)
            Degree of estimated drift polynomial.
        sig0_mugal : int, optional (default=1)
            A priori standard deviation of unit weight of observations [µGal] for the stochastic model of the
            least-squares adjustment.
        scaling_factor_datum_observations : float, optional (default=1.0)
            Factor for scaling the standard deviation (SD) of g of datum stations. The scaled SD is is used for
            weighting the direct pseudo observations of g at the datum stations that are introduced as datum
            constraints.
        add_const_to_sd_of_observations_mugal : float, optional (default=0.0)
            The defined additive constant is added to the standard deviation (SD) of setup
            observations in order to scale the SD and the resulting weights to realistic values. In µGal. The scaling
            factor `scaling_factor_for_sd_of_observations` is applied before adding this constant!
        scaling_factor_for_sd_of_observations : float, optional (default=1.0)
            Scaling factor for the standard deviation of the setup observations. `add_const_to_sd_of_observations_mugal`
            is applied after applying the scaling factor!
        confidence_level_chi_test : float, optional (default=0.95)
            Confidence level for the goodness-of-fit test.
        confidence_level_tau_test : float, optional (default=0.95)
            Confidence level for the tau test.
        drift_ref_epoch_type : string ('survey' or 'campaign'), optional (default='survey')
            Defines whether the reference epoch t0 for the estimation of the drift polynomials for each survey in the
            campaign is the reference epoch of the first (active) observation in each survey (option: 'survey') or the
            first (active) observation in the whole campaign (option: 'campaign').
        noise_floor_mugal : float, optional (default=0.0)
            The standard error SE of the estimated gravity values at stations is calculated by SE = sqrt(SD**2 + NF**2),
            where SD is the estimated standard deviation of the station's gravity and NF is the noise floor value
            specified here.
        verbose : bool, optional (default=False)
            If `True`, status messages are printed to the command line, e.g. for debugging and testing
        """

        # Init.:
        flag_s0_within_threshold = False
        # scale_factor = scaling_factor_for_sd_of_observations
        add_const = add_const_to_sd_of_observations_mugal  # initial
        add_const_step_size = add_const_to_sd_of_observations_step_size_mugal  # [mugal]
        mult_factor = scaling_factor_for_sd_of_observations  # initial
        mult_factor_step_size = multiplicative_factor_step_size_percent / 100  # [%] => factor
        last_iteration_step_action = ''  # 'increase' or 'decrease'
        iteration_log_str = ''
        complete_log_str = ''
        add_const_total_mugal = 0
        mult_factor_total = 1.0

        # Run iteration:
        for i_iteration in range(1, max_number_iterations + 1):

            # Only apply multiplicative facot or additive constant in the first iteration, depending on the iteration
            # approach:
            if i_iteration > 1:
                if iteration_approach == 'Additive':
                    mult_factor = 1.0
                elif iteration_approach == 'Multiplicative':
                    add_const = 0.0

            add_const_total_mugal += add_const  # Log total additive constant
            mult_factor_total *= mult_factor  # Log total multiplicative factor
            self.log_str = ''  # Reset log string of the previous run

            self.adjust(drift_pol_degree=drift_pol_degree,
                        sig0_mugal=sig0_mugal,
                        scaling_factor_datum_observations=scaling_factor_datum_observations,
                        add_const_to_sd_of_observations_mugal=add_const,  # Adjusted iteratively
                        scaling_factor_for_sd_of_observations=mult_factor,  # Adjusted iteratively
                        confidence_level_chi_test=confidence_level_chi_test,
                        confidence_level_tau_test=confidence_level_tau_test,
                        drift_ref_epoch_type=drift_ref_epoch_type,
                        noise_floor_mugal=noise_floor_mugal,
                        verbose=False
                        )

            iteration_log_str_tmp = f'########## Iteration {i_iteration} #########################\n'
            # if iteration_approach == 'Additive':
            iteration_log_str_tmp += f'Total additive const. to SD: {add_const_total_mugal:5.3f} µGal\n'
            iteration_log_str_tmp += f'Current additive const. to SD: {add_const:5.3f} µGal\n'
            # elif iteration_approach == 'Multiplicative':
            iteration_log_str_tmp += f'Total mult. factor for SD: {mult_factor_total:5.3f}\n'
            iteration_log_str_tmp += f'Current mult. factor for SD: {mult_factor:5.3f}\n'
            iteration_log_str_tmp += f's0² a posteriori: {self.s02_a_posteriori:5.3f} µGal²\n'
            iteration_log_str_tmp += f'\n'
            if verbose:
                print(iteration_log_str_tmp)
                print(self.log_str)
            if self.write_log:
                complete_log_str = complete_log_str + iteration_log_str_tmp + self.log_str + '\n\n'
            iteration_log_str += iteration_log_str_tmp

            if (self.s02_a_posteriori < (s02_target + s02_target_delta)) and (
                    self.s02_a_posteriori > (s02_target - s02_target_delta)):
                flag_s0_within_threshold = True
            elif self.s02_a_posteriori > (s02_target + s02_target_delta):  # Too large => Increase SD of obs.
                if last_iteration_step_action == 'decrease':
                    add_const_step_size = add_const_step_size / 2
                    mult_factor_step_size = mult_factor_step_size / 2
                last_iteration_step_action = 'increase'
                add_const = add_const_step_size
                mult_factor = 1.0 + mult_factor_step_size
            elif self.s02_a_posteriori < (s02_target - s02_target_delta):  # Too small => Decrease SD of obs.
                if last_iteration_step_action == 'increase':
                    add_const_step_size = add_const_step_size / 2
                    mult_factor_step_size = mult_factor_step_size / 2
                last_iteration_step_action = 'decrease'
                add_const = -add_const_step_size
                mult_factor = 1.0 - mult_factor_step_size

            if flag_s0_within_threshold:  # Exit loop and stop iteration
                break

        # Write status message to log according to iteration result:
        if iteration_approach == 'Additive':
            if flag_s0_within_threshold and (add_const_total_mugal <= max_total_additive_const_to_sd_mugal):
                iteration_log_str_tmp = f' => Iteration successful!\n'
                iteration_log_str_tmp += f' => s0² a posteriori of {self.s02_a_posteriori:1.3f} within [{s02_target - s02_target_delta:1.3f}, {s02_target + s02_target_delta:1.3f}]\n'
                iteration_log_str_tmp += f' => Total additive constant to SD of observations ({add_const_total_mugal:1.3f}) ' + \
                                         f'is smaller than the user defined threshold ' + \
                                         f'of {max_total_additive_const_to_sd_mugal:1.3f} µGal.\n'
            else:
                iteration_log_str_tmp = f' => ERROR: Iteration failed!\n'
                if not flag_s0_within_threshold:
                    iteration_log_str_tmp += f' => s0² a posteriori of {self.s02_a_posteriori:1.3f} not within ' + \
                                             f'[{s02_target - s02_target_delta:1.3f}, {s02_target + s02_target_delta:1.3f}]\n'
                if (add_const_total_mugal > max_total_additive_const_to_sd_mugal):
                    iteration_log_str_tmp += f' => Total additive constant to SD  ({add_const_total_mugal:1.3f}) ' + \
                                             f'exceeds the the user defined threshold ' + \
                                             f'of {max_total_additive_const_to_sd_mugal:1.3f} µGal.\n'
            iteration_log_str_tmp += f'\n'

        elif iteration_approach == 'Multiplicative':
            if flag_s0_within_threshold and (mult_factor_total <= (max_multiplicative_factor_to_sd_percent/100)) and (mult_factor_total >= (min_multiplicative_factor_to_sd_percent/100)):
                iteration_log_str_tmp = f' => Iteration successful!\n'
                iteration_log_str_tmp += f' => s0² a posteriori of {self.s02_a_posteriori:1.3f} within [{s02_target - s02_target_delta:1.3f}, {s02_target + s02_target_delta:1.3f}]\n '
                iteration_log_str_tmp += f' => Total multiplicative factor for SD of observations ({mult_factor_total*100:1.3f}%) ' + \
                                         f'is between the user defined threshold ' + \
                                         f'of {min_multiplicative_factor_to_sd_percent:1.3f}% and {max_multiplicative_factor_to_sd_percent:1.3f}%.\n'
            else:
                iteration_log_str_tmp = f' => ERROR: Iteration failed!\n'
                if not flag_s0_within_threshold:
                    iteration_log_str_tmp += f' => s0² a posteriori of {self.s02_a_posteriori:1.3f} not within ' + \
                                             f'[{s02_target - s02_target_delta:1.3f}, {s02_target + s02_target_delta:1.3f}]\n'
                if (mult_factor_total > (max_multiplicative_factor_to_sd_percent/100)):
                    iteration_log_str_tmp += f' => Total multiplicative factor for SD  ({mult_factor_total*100:1.3f}%) ' + \
                                             f'exceeds the the user defined threshold ' + \
                                             f'of {max_multiplicative_factor_to_sd_percent:1.3f}%.\n'
                if (mult_factor_total < (min_multiplicative_factor_to_sd_percent/100)):
                    iteration_log_str_tmp += f' => Total multiplicative factor for SD  ({mult_factor_total*100:1.3f}%) ' + \
                                             f'is smaller than the user defined threshold ' + \
                                             f'of {min_multiplicative_factor_to_sd_percent:1.3f}%.\n'
            iteration_log_str_tmp += f'\n'

        iteration_log_str += iteration_log_str_tmp

        if verbose:
            print(iteration_log_str_tmp)

        # Append iteration log to the log string of the last iteration:
        if self.write_log:
            self.log_str = complete_log_str + \
                           '\n########## Iteration log #########################\n\n' + \
                           f' - Iteration approach: {iteration_approach}\n\n' + \
                           iteration_log_str

        if not flag_s0_within_threshold:
            raise AssertionError(f'Iteration Error: s0² a posteriori of {self.s02_a_posteriori:1.3f} not within '
                                 f'[{s02_target - s02_target_delta:1.3f}, {s02_target + s02_target_delta:1.3f}] after '
                                 f'{i_iteration} iterations!')

        if iteration_approach == 'Additive':
            if abs(add_const_total_mugal) > max_total_additive_const_to_sd_mugal:
                raise AssertionError(f'Total additive constant to SD  ({add_const_total_mugal:1.3f}) '
                                     f'exceeds the the user defined threshold '
                                     f'of {max_total_additive_const_to_sd_mugal:1.3f} µGal.\n')
        elif iteration_approach == 'Multiplicative':
            if mult_factor_total > (max_multiplicative_factor_to_sd_percent/100):
                raise AssertionError(f'Total multiplicative factor for SD  ({mult_factor_total*100:1.3f}%) '
                                     f'exceeds the the user defined threshold '
                                     f'of {max_multiplicative_factor_to_sd_percent:1.3f}%.\n')
            if mult_factor_total < (min_multiplicative_factor_to_sd_percent/100):
                raise AssertionError(f'Total multiplicative factor for SD  ({mult_factor_total*100:1.3f}%) '
                                     f'is smaller than the the user defined threshold '
                                     f'of {min_multiplicative_factor_to_sd_percent:1.3f}%.\n')

        self.number_of_iterations = i_iteration

    def export_stat_results_shapefile(self, filename, epsg_code, verbose=True):
        """Export station-related results from the `stat_obs_df` dataframe to a shapefile.

        Notes
        -----
        This method relies on the optional module `geopandas` (optional dependency).

        Parameters
        ---------
        filename : str
            Name and path of the output shapefile.
        epsg_code: int
            EPSG code of the station coordinates' CRS.
        verbose : bool, optional (default=False)
            `True` implies that status messages are printed to the command line.
        """
        if not _has_geopandas:
            raise ImportError(f'Optional dependency geopandas not available, but needed for writing shapefiles!')
        if verbose:
            print(f'Save station results of lsm run "{self.comment}" to: {filename}')

        stat_obs_df = self.stat_obs_df.copy(deep=True)
        stat_obs_df['comment'] = self.comment
        stat_obs_gdf = geopandas.GeoDataFrame(stat_obs_df, geometry=geopandas.points_from_xy(stat_obs_df['long_deg'], stat_obs_df['lat_deg']), crs=epsg_code)
        stat_obs_gdf.to_file(filename)

    @property
    def time_str(self):
        """Return time of lsm adjustment as formatted string."""
        return self.init_time.strftime("%Y-%m-%d, %H:%M:%S")

    @property
    def get_log_string(self):
        """Returns the log string."""
        return self.log_str

    @property
    @time_it
    def get_correlation_matrix(self):
        """Calculates und returns the correlation matrix based on the Co-Variance matrix of the LSM run.

        Returns
        -------
        numpy array : Correlation matrix (n, n)

        Notes
        -----
        The co-variance matrix `self.Cxx` needs to exist.
        """
        if self.Cxx is None:
            # raise AssertionError('Correlation matrix cannot be calculated due to missing co-variance matrix.')
            return None
        mat_Rxx = np.full((self.Cxx.shape[0], self.Cxx.shape[1]), np.nan)
        for i_row in range(0, mat_Rxx.shape[0]):
            cxx_row_sqrt = np.sqrt(self.Cxx[i_row, i_row])
            for i_col in range(0, i_row+1):
                mat_Rxx[i_row, i_col] = self.Cxx[i_row, i_col] / (
                            cxx_row_sqrt * np.sqrt(self.Cxx[i_col, i_col]))
                mat_Rxx[i_col, i_row] = mat_Rxx[i_row, i_col]
        return mat_Rxx

    @property
    def get_results_obs_df(self):
        """Getter for the observation-related results."""
        return self.setup_obs_df

    @property
    def get_results_drift_df(self):
        """Getter for the drift-related results."""
        return self.drift_pol_df

    @property
    def get_results_stat_df(self):
        """Getter for the station-related results."""
        return self.stat_obs_df

    @property
    def get_results_vg_df(self):
        """Getter for the results of vertical gravity gradient estimation."""
        try:
            return self.vg_pol_df
        except AttributeError:  # If "self" has no attribute "vg_pol_df" (not initialized)
            return None

    @property
    def check_for_unique_setup_id(self):
        """Indicates whether setup IDs should be unique based on the setup data calculation method.

        Returns
        -------
        bool : `True`, if the setup IDs should be unique, otherwise `False`.

        Notes
        -----
        If observations were not aggregated within at least one setup, checking for unique setup IDs makes no sense.
        """
        for survey_name, setup_data in self.setups.items():
            if setup_data['setup_calc_method'] == 'individual_obs':
                return False
        return True

    @property
    def survey_info_string(self) -> str:
        """Returns information on all surveys."""
        survey_names_list = []
        num_setups_list = []
        num_obs_total_list = []
        num_obs_active_list = []
        tide_correction_type_list = []
        reference_height_type_list = []
        atm_pres_correction_type_list = []
        setup_calc_method_list = []
        for survey_name, survey in self.setups.items():
            survey_names_list.append(survey_name)
            setup_df = survey['setup_df']
            num_setups_list.append(len(setup_df))
            setup_obs_list_df = survey['setup_obs_list_df']
            num_obs_total_list.append(len(setup_obs_list_df))
            num_obs_active_list.append(len(setup_obs_list_df.loc[setup_obs_list_df['keep_obs']]))
            tide_correction_type_list.append(survey['tide_correction_type'])
            reference_height_type_list.append(survey['reference_height_type'])
            atm_pres_correction_type_list.append(survey['atm_pres_correction_type'])
            setup_calc_method_list.append(survey['setup_calc_method'])
        tmp_df = pd.DataFrame(list(zip(survey_names_list,
                                       num_setups_list,
                                       num_obs_total_list,
                                       num_obs_active_list,
                                       tide_correction_type_list,
                                       reference_height_type_list,
                                       atm_pres_correction_type_list,
                                       setup_calc_method_list,
                                       )),
                          columns=['Survey', '#Setups', '#Obs', '#Active obs', 'Tide correction', 'Ref. height', 'Atm. correction', 'Setup calc. method'])
        return tmp_df.to_string(index=False)

    # @property
    # def obs_level_corrections_info_string(self) -> str:
    #     """Returns information on the applied observation-level corrections."""
    #     info_str = ''
    #     return info_str

def bin_redundacy_components(mat_r):
    """Bin redundancy components for easier interpretatzion.

    See: Skriptum AG2 (Navratil, TU Wien), p. 70.
    """
    number_obs_r_equal_0 = 0  # No error control: Errors cannot be detected
    number_obs_r_0_to_03 = 0  # Bad error control
    number_obs_r_03_to_07 = 0  # Good error control
    number_obs_r_07_1 = 0  # Very good error control, but observation may be redundant
    number_obs_r_equal_1 = 0  # Observatiuon is redundant

    results_dict = {}  # '': ''

    return number_obs_r_equal_0, number_obs_r_0_to_03

def tau_test(mat_w, dof, alpha, mat_r):
    """Tau-criterion test for outlier detection of a least-squares adjustment.

    This test considers that the true variance factor sigma_0² is unknown. It is based on normalized residuals
    (Studentized residuals) and the Tau distribution.

    If the test fails it indicates that ONE (!) observation is a gross error, i.e. ONE residual is an outlier!
    If the adjustment is affected by two or more gross errors this test is not applicable. In that case the
    following pragmatic approach is recommended (e.g. by Caspary, 1987): The observation with the largest test
    statistic is discarded. Then the adjustment is repeated with the remaining n-1 observations and the test is
    repeated based on the new results, etc.

    Notes
    -----
    Ref.: Pope (1976): The statistics of residuals and the detection of outliers, NOAA Technical Report NOS 65 NGS 1
    Ref.: Caspari (1987): Concepts of network and deformation analysis. Monograph by UNSW Sydney, pp. 76-77
    This method implements a two-sided test, as the critical value is calculated with (alpha/2) and the abs. value of
    the normalized residual w is tested.

    Parameters
    ----------
    mat_w: np.array of floats
        Vector with standardized post-fit residuals of all observations in the least-squares adjustment.
        They are computed as post-fit residuals divides by their standard deviations.
    dof: int
        Degree of freedom of the least-squares adjustment.
    alpha: float (0 to 1)
        Significance level.
    mat_r: np.array of floats
        Vector with redundancy components derived in the least-squares adjustment for each observation.
    """
    # Critical value (two-sided test: the abs. value of w is tested und the Tau-distribution is symmetric about 0!):
    tsd_crt = stats.t.ppf(1 - alpha / 2, dof - 1)  # t distribution
    tau_crt = (tsd_crt * np.sqrt(dof)) / np.sqrt(dof - 1 + tsd_crt ** 2)  # Critical value (Pope, 1976, Eq. (6))
    tau_test_result = []
    for idx, w in enumerate(mat_w):
        # Check, if redancy component is larger than threshold:
        if mat_r[idx] > settings.R_POPE_TEST_TRESHOLD:
            if np.abs(w) > tau_crt:  # Equ. (6-36) in Caspary (1987)
                tau_test_result.append('failed')
            else:
                tau_test_result.append('passed')
        else:
            tau_test_result.append(f'r too small')  # Pope test not applied due to small redundancy component!
    return tau_test_result, tau_crt

def create_hist(mat_v):
    """Create histogram."""
    residuals = np.ndarray.tolist(mat_v)
    residuals = [item for sublist in residuals for item in sublist]
    hist_residuals, bin_edges = np.histogram(residuals, bins=5)
    return hist_residuals, bin_edges

def global_model_test(cf, dof, a_posteriori_variance_of_unit_weight, a_priori_variance_of_unit_weight):
    """Global model test based on the comparison of a posteriori and a priori variance of unit weight.

    This global model test asserts if the "model is correct and complete". If it fails it indicates that the
    observations contradict the mathematical adjustment model. However, this test cannot the validity of the model or
    the correctness of the observation!

    Notes
    -----
    This "global model test" is described by Caspari (1987): Concepts of network and deformation analysis. Monograph
    by UNSW Sydney, pp. 6-8 and pp.68-69.
    This method implements a two-sided test.

    Parameters
    ----------
    cf: float
        Confidence level for Chi²
    dof: int
        Degree of freedom.
    a_posteriori_variance_of_unit_weight: float
        A posteriori variance of unit weight (after adjustment).
    a_priori_variance_of_unit_weight: float
        A priori variance of unit weight (before adjustment).

    Returns
    -------
    """
    alpha = 1 - cf  # Significance level = Probability of committing a type 1 error (H0 wrongly dismissed)
    chi_crit_upper = stats.chi2.ppf(1 - alpha / 2, dof)  # critical value
    chi_crit_lower = stats.chi2.ppf(alpha / 2, dof)  # critical value
    # chi_critical_value = stats.chi2.ppf(alpha, dof)  # critical value
    test_value = dof * a_posteriori_variance_of_unit_weight / a_priori_variance_of_unit_weight  # Equ. (2-13) in Caspary (1987)
    if chi_crit_lower < test_value < chi_crit_upper:
    # if test_value <= test_value:  # Equ. (2-15) in Caspary (1987)
        chi_test_status = 'Passed'
    else:
        chi_test_status = 'Not passed'
    chi_crit = [chi_crit_lower, chi_crit_upper]
    return chi_crit, test_value, chi_test_status


from fastapi import FastAPI

from src.db import models
from src.db.database import engine
from src.recipes import recipe_routers
from src.users import user_routers

description = """
Recipes_API helps you do awesome stuff. ðŸš€
## Recipes
You can:
* create recipe
* read recipe
* update recipe
* delete recipe
* like recipe
* add recipe to favorite
* remove recipe from favorite
* block recipe (if you are admin)
* unblock recipe (if you are admin)

## Users
You can:
* create user
* authenticate user
* read user
* update user
* delete user
* block user (if you are admin)
* unblock user (if you are admin)

"""

app = FastAPI(
    title="FastAPI_Recipes",
    description=description,
    version="0.0.1",

)

models.Base.metadata.create_all(bind=engine)

app.include_router(user_routers.router_auth)
app.include_router(user_routers.router)
app.include_router(recipe_routers.router)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run('main:app', host="localhost", port=8000, reload=True)


############################
#                          #
#将两个表格的数据读入并且合并#
#                          #
############################
import xlrd
import xlwt

#读取Excel数据
def GetExcel(dir):
    data = xlrd.open_workbook(dir)
    table = data.sheet_by_name("Sheet1")
    rows = table.nrows
    col = table.ncols
    DataExcel = []
    dataExcel = {}
    for i in range(rows):
        DataExcel.append(table.row_values(i))
    for i in range(1, len(DataExcel)):
        n1 = int(DataExcel[i][0]) + 202000
        DataExcel[i][0] = str(n1)
        if(DataExcel[i][-1]=='bad'):
            DataExcel[i][-1]='1'
        elif(DataExcel[i][-1]=='general'):
            DataExcel[i][-1]='2'
        elif(DataExcel[i][-1]=='good'):
            DataExcel[i][-1]='3'
        elif(DataExcel[i][-1]=='excellent'):
            DataExcel[i][-1]='4'
        x = {DataExcel[i][0]: DataExcel[i][1:]}
        dataExcel.update(x)
    return dataExcel

#读取txt中数据
def GetTxt(dir):
    file = open(dir)
    dataTxt = []
    DataTxt = {}
    while 1:
        line = file.readline()
        line = line[:-1]
        if not line:
            break
        line = line.split(",")
        if (line[-1] == 'bad'):
            line[-1] = '1'
        elif (line[-1] == 'general'):
            line[-1] = '2'
        elif (line[-1] == 'good'):
            line[-1] = '3'
        elif (line[-1] == 'excellent'):
            line[-1] = '4'
        if(not line[0] in DataTxt.keys()):
            x = {line[0]:line[1:]}
            DataTxt.update(x)
        else:
            for i in range(1,16):
                if(DataTxt[line[0]][i-1] ==''):
                    DataTxt[line[0]][i-1] = line[i]
    file.close()
    return DataTxt

#合并TXT以及Execl数据
def dataMerge():
    dataExcel = GetExcel("D:/11.xlsx")
    dataTxt = GetTxt("D:/11.txt")
    dataALl = {}
    dataALl = dataTxt;
    for key in dataALl.keys():
        for i in range(0,15):
          if(not i == 13):
                if(key in dataExcel.keys()):
                   if(dataALl[key][i]==''):
                        if(not dataExcel[key][i] =='' ):
                            dataALl[key][i] = dataExcel[key][i]
                    #dataALl.update("{{}:{}}".format(key,dataExcel[key]))
    for key in dataExcel.keys():
        if(not key in dataALl):
            x = "{'" + "{}':{}".format(key, dataExcel[key])+"}"
            x = eval(x)
            dataALl.update(x)
    for key in dataTxt.keys():
        if(not key in dataALl):
            dataALl.update(dataTxt[key])

    for key in dataALl.keys():
        for i in range(0,15):
            if(not i == 13):
                if (key in dataTxt.keys()):
                    if(dataALl[key][i]==''):
                        if(not dataTxt[key][i] ==''):
                            dataALl[key][i] = dataTxt[key][i]
    return  dataALl
#将合并的数据写入excel
'''
def writeExcel(dataALl):
    workbook = xlwt.Workbook(encoding = 'utf-8')

    worksheet = workbook.add_sheet('My Worksheet')

    i = 1;
    for key in dataALl.keys():
        worksheet.write(i, 0, label=key)
        for k in range(15):
            worksheet.write(i, k+1, label=dataALl[key][k])
        i = i+1

    workbook.save('Excel_test.xls')
'''


import subprocess
import argparse
import numpy as np
from PIL import Image
from matplotlib import cm
import pickle
from transformers import BartTokenizer, BartTokenizerFast
import torch
from transformers import AutoFeatureExtractor
from torchvision import transforms
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch.nn as nn
import sacrebleu as scb
# from moverscore_v2 import get_idf_dict, word_mover_score
from collections import defaultdict

def eval_totto(prediction_path, target_path):
    command = 'bash ../language/totto/totto_eval.sh --prediction_path ' + prediction_path \
              + ' --target_path ' + target_path
    # command = 'bash ' + eval_script_path+' --prediction_path ' + prediction_path \
    #           + ' --target_path ' + target_path
    try:
        result = subprocess.run(command,
                                check=True,
                                shell=True,
                                stdout=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(
            "command '{}' return with error (code {}): {}".format(e.cmd, e.returncode, e.output))
    res = result.stdout.decode("utf-8")
    content_list = res.split(r'BLEU+case.mixed+numrefs.3+smooth.exp+tok.13a+version.1.4.10 = ')
    overall_bleu = float(content_list[1].split()[0])
    overlap_bleu = float(content_list[2].split()[0])
    nonoverlap_bleu = float(content_list[3].split()[0])
    return overall_bleu, overlap_bleu, nonoverlap_bleu

def parse_config():
    parser = argparse.ArgumentParser()
    # data configuration
    parser.add_argument('--save_model_name', type=str)
    parser.add_argument('--save_output', type=str)
    parser.add_argument('--output_name', type=str)
    parser.add_argument('--process_mode', type=str, default="AVG_POOLING_LAYER")
    parser.add_argument('--train_table_text_path', type=str, default="../data/train/totto_reasoning_table.txt")
    parser.add_argument('--train_table_header_path', type=str, default='')
    parser.add_argument('--train_content_text_path', type=str, default='../data/train/totto_reasoning_used_headers.txt')
    parser.add_argument('--train_table_highlight_path', type=str, default='../data/train/totto_reasoning_highlight_headers.txt')
    parser.add_argument('--train_reference_sentence_path', type=str, default='../data/train/totto_reasoning_reference.txt')
    parser.add_argument('--train_path_text_path', type=str)
    parser.add_argument('--test_table_text_path', type=str)
    parser.add_argument('--test_table_header_path', type=str)
    parser.add_argument('--test_content_text_path', type=str)
    parser.add_argument('--test_table_highlight_path', type=str)
    parser.add_argument('--test_reference_sentence_path', type=str)
    parser.add_argument('--test_path_text_path', type=str)
    parser.add_argument('--dev_table_text_path', type=str, default='../data/dev/totto_reasoning_table.txt')
    parser.add_argument('--dev_table_header_path', type=str)
    parser.add_argument('--dev_content_text_path', type=str, default='../data/dev/totto_reasoning_used_headers.txt')
    parser.add_argument('--dev_reference_sentence_path', type=str, default='../data/dev/totto_reasoning_reference.txt')
    parser.add_argument('--dev_path_text_path', type=str)
    parser.add_argument('--dev_table_highlight_path', type=str, default='../data/dev/totto_reasoning_highlight_headers.txt')
    parser.add_argument('--batch_size', type=int, default=48)
    parser.add_argument('--pretrained_ckpt_path', type=str)
    parser.add_argument('--planner_pretrained_ckpt_path', type=str)
    parser.add_argument('--finetune_generator_ckpt_path', type=str)
    parser.add_argument('--backbone_generator_ckpt_path', type=str)
    parser.add_argument('--special_token_path', type=str, default='../DataAnalysis/totto_col_header_vocab.txt')
    parser.add_argument('--dev_reference_path', type=str, default='../data/raw_data/totto_dev_data.jsonl')
    parser.add_argument('--max_table_len', type=int, default=256)
    parser.add_argument('--max_content_plan_len', type=int, default=128)
    parser.add_argument('--max_tgt_len', type=int, default=256)
    parser.add_argument('--min_slot_key_cnt', type=int, default=10)
    parser.add_argument('--generator_model_name', type=str, default='facebook/bart-base')
    parser.add_argument('--planner_model_name', type=str, default='microsoft/resnet-50')
    parser.add_argument('--max_decode_len', type=int, default=256)
    parser.add_argument('--gpu_id', type=int, default=0)
    parser.add_argument('--total_steps', type=int, default=200000)
    parser.add_argument('--warmup_steps', type=int, default=2000)
    parser.add_argument('--learning_rate', type=float, default=2e-5)
    parser.add_argument('--print_every', type=int, default=200)
    parser.add_argument('--eval_every', type=int, default=2000)
    parser.add_argument('--gradient_accumulation_steps', type=int, default=2)
    parser.add_argument('--dropout', type=float, default=0.2)
    parser.add_argument('--ckpt_path', type=str, default=r'./ckpt/bart_pretrain/')
    parser.add_argument('--inference_output_file', type=str)

    return parser.parse_args()

def map_cuda(tensor_item, device, is_cuda):
    res_list = []
    if is_cuda:
        res_list.append(tensor_item[0].cuda(device))
        res_list.append(tensor_item[1].cuda(device))
    else:
        res_list = tensor_item
    return res_list

def map_cuda_(tensor_list, device, is_cuda):
    if is_cuda:
        res_list = tensor_list.cuda(device)
    else:
        res_list = tensor_list
    return res_list

def pickle_write(file_path, data):
    f = open(file_path, 'wb')
    pickle.dump(data, f)
    f.close()

def pickle_read(file_path):
    f = open(file_path, 'rb')
    data = pickle.load(f)
    return data

def process_attention_data_extractor(one_cross_attention):
    # print("process_attention_data")
    extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
    one_input = extractor(one_cross_attention, return_tensors='pt')[
        'pixel_values']
    # print(one_input)
    return one_input

def process_attention_data(one_cross_attention):
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    # print("process")
    one_input_tensor = preprocess(one_cross_attention).unsqueeze(0)
    # print(one_input_tensor)
    return one_input_tensor

def img_evaluate(predicted_labels, labels):
    # print(len(predicted_labels), len(labels))
    assert (len(predicted_labels) == len(labels))
    print(type(predicted_labels))
    print(type(labels))
    tp = (labels * predicted_labels).sum().to(torch.float32)
    tn = ((1 - labels) * (1 - predicted_labels)).sum().to(torch.float32)
    fp = ((1 - labels) * predicted_labels).sum().to(torch.float32)
    fn = (labels * (1 - predicted_labels)).sum().to(torch.float32)

    # print(tp, tn, fp, fn)
    precision = tp / (tp + fp)
    accuracy = (tp + tn) / (tp + tn + fp + fn)
    recall_score = tp / (tp + fn)
    f1_score = 2 * (precision * recall_score) / (precision + recall_score)

    return precision, accuracy, recall_score, f1_score

def cvtOneTensorToPIL(inputTensor):
    # print(inputTensor.device)
    input_tensor = inputTensor
    im = Image.fromarray(np.uint8(cm.gist_earth(input_tensor) * 255), mode="RGB")
    return im

def initial_tokenizer(special_token_path="../../data/totto_col_header_vocab.txt",
                      min_slot_key_cnt=10):
    model_name = "facebook/bart-base"
    special_token_list, special_token_dict = [], {}
    with open(special_token_path, 'r', encoding='utf8') as i:
        lines = i.readlines()
        for l in lines:
            one_special_token = l.strip('\n').split()[0]
            cnt = int(l.strip('\n').split()[1])
            if cnt >= min_slot_key_cnt:
                special_token_list.append(one_special_token)
                special_token_dict[one_special_token] = 1
            else:
                pass
    print('Number of Special Token is %d' % len(special_token_list))
    tokenizer = BartTokenizerFast.from_pretrained(model_name)
    decode_tokenizer = BartTokenizer.from_pretrained(model_name)
    print('original vocabulary Size %d' % len(tokenizer))
    tokenizer.add_tokens(special_token_list)
    decode_tokenizer.add_tokens(special_token_list)
    print('vocabulary size after extension is %d' % len(tokenizer))
    return tokenizer, decode_tokenizer

def process_header_data(tokenizer, batch_size, batch_table_id_list, batch_header_id_list,
                        batch_highlight_id_list, batch_encoder_attentions, batch_cross_attentions):
    batch_header_cross_attention_list = []
    batch_header_cat_attention_list = []
    batch_labels = []
    batch_header_dict = {}
    batch_encoder_attentions = batch_encoder_attentions[-1]
    batch_cross_attentions = batch_cross_attentions[-1]

    for bitem in range(batch_size):

        src_text_id_list = batch_table_id_list[bitem]
        header_text_id_list = batch_header_id_list[bitem]
        highlight_text_id_list = batch_highlight_id_list[bitem]

        rm_special_token_ids = [tokenizer.convert_tokens_to_ids('<pad>'),
                                tokenizer.convert_tokens_to_ids('__SEP__'),
                                tokenizer.convert_tokens_to_ids('<pad>'),
                                tokenizer.convert_tokens_to_ids('__None__'),
                                tokenizer.convert_tokens_to_ids('__None__'),
                                tokenizer.convert_tokens_to_ids('__EOS__'),
                                tokenizer.convert_tokens_to_ids('__PAD__'),
                                1437]

        # all headers
        header_text_id_set = set(header_text_id_list)

        for special_token in rm_special_token_ids:
            if special_token in header_text_id_set:
                header_text_id_set.remove(special_token)

        # get header position in src text
        position = []
        for i, src_token_id in enumerate(src_text_id_list):
            if src_token_id in header_text_id_set:
                position.append(i)

        cross_attention = batch_cross_attentions[bitem]
        encoder_attention = batch_encoder_attentions[bitem]

        header_cross_attention_list = []
        header_cat_attention_list = []
        labels = []

        # extract all headers' attentions & labels
        for i in position:
            if (src_text_id_list[i]) in highlight_text_id_list:
                labels.append(1)
            else:
                labels.append(0)

            header_cross_attention_i = cross_attention[:, :, i]
            header_encoder_attention_i = encoder_attention[:, :, i]

            header_cross_attention_PIL = cvtOneTensorToPIL(header_cross_attention_i.detach())
            header_cross_attention_list.append(header_cross_attention_PIL)

            cat_attention = torch.cat((header_encoder_attention_i, header_cross_attention_i),
                                      dim=1)
            cat_attention_PIL = cvtOneTensorToPIL(cat_attention.detach())
            header_cat_attention_list.append(cat_attention_PIL)

        batch_header_cross_attention_list.append(header_cross_attention_list)
        batch_header_cat_attention_list.append(header_cat_attention_list)
        batch_labels.append(labels)

    batch_header_dict['crossAttentions'] = batch_header_cross_attention_list
    batch_header_dict['catAttentions'] = batch_header_cat_attention_list
    batch_header_dict['labels'] = batch_labels
    return batch_header_dict

def process_sentence_header_data(tokenizer, data_table_id_list, data_header_id_list,
                                 data_highlight_id_list, data_encoder_attentions,
                                 data_cross_attentions):

    data_labels = []
    data_header_dict = {}

    ln_encoder_attentions = []
    for i in data_encoder_attentions:
        i = i.cpu()
        layer_norm_i = nn.LayerNorm([i.shape[1], i.shape[2], i.shape[3]])
        data_encoder_attention = layer_norm_i(i)
        ln_encoder_attentions.append(data_encoder_attention)

    ln_cross_attentions = []
    for i in data_cross_attentions:
        i = i.cpu()
        layer_norm_i = nn.LayerNorm([i.shape[1], i.shape[2], i.shape[3]])
        data_cross_attention = layer_norm_i(i)
        ln_cross_attentions.append(data_cross_attention)

    src_text_id_list = data_table_id_list
    header_text_id_list = data_header_id_list
    highlight_text_id_list = data_highlight_id_list

    print("len highlight id list: ", len(highlight_text_id_list))
    # print("header id list: ", len(header_text_id_list))

    rm_special_token_ids = [tokenizer.convert_tokens_to_ids('<pad>'),
                            tokenizer.convert_tokens_to_ids('__SEP__'),
                            tokenizer.convert_tokens_to_ids('<pad>'),
                            tokenizer.convert_tokens_to_ids('__None__'),
                            tokenizer.convert_tokens_to_ids('__None__'),
                            tokenizer.convert_tokens_to_ids('__EOS__'),
                            tokenizer.convert_tokens_to_ids('__PAD__'),
                            1437]

    # all headers
    header_text_id_set = set(header_text_id_list)

    # print(rm_special_token_ids)
    for special_token in rm_special_token_ids:
        if special_token in header_text_id_set:
            header_text_id_set.remove(special_token)

    # print("len set highlight id list: ", len(header_text_id_set))
    # get header position in src text
    position = []
    for i, src_token_id in enumerate(src_text_id_list):
        if src_token_id in header_text_id_set:
            position.append(i)

    header_cross_attention_list = []
    header_cat_attention_list = []
    header_tokens = []

    # extract all headers' attentions & labels
    for i in position:
        if (src_text_id_list[i]) in highlight_text_id_list:
            data_labels.append(1)
        else:
            data_labels.append(0)

        header_token = tokenizer.convert_ids_to_tokens(src_text_id_list[i])
        header_tokens.append(header_token)

        encoder_attentions = []
        for encoder_attention in ln_encoder_attentions:
            encoder_attention = encoder_attention.squeeze(0)
            header_encoder_attention_i = encoder_attention[:, :, i]
            encoder_attentions.append(header_encoder_attention_i)

        cross_attentions = []
        for cross_attention in ln_cross_attentions:
            cross_attention = cross_attention.squeeze(0)
            header_cross_attention_i = cross_attention[:, :, i]
            cross_attentions.append(header_cross_attention_i)

        cat_encoder_attentions = encoder_attentions[0]
        for encoder_a in encoder_attentions[1:]:
            cat_encoder_attentions = torch.cat([cat_encoder_attentions, encoder_a],
                                               dim=0)

        cat_cross_attentions = cross_attentions[0]
        for cross_a in cross_attentions[1:]:
            cat_cross_attentions = torch.cat([cat_cross_attentions, cross_a], dim=0)

        header_cross_attention_PIL = cvtOneTensorToPIL(cat_cross_attentions)
        header_cross_attention_list.append(header_cross_attention_PIL)

        cat_attention = torch.cat((cat_encoder_attentions, cat_cross_attentions),
                                                            dim=1)

        cat_attention_PIL = cvtOneTensorToPIL(cat_attention)
        header_cat_attention_list.append(cat_attention_PIL)

        # print("cat_encoder_attentions shape ", cat_encoder_attentions.shape)
        # print("cat_cross_attentions shape ", cat_cross_attentions.shape)
        # print("cat_attention shape ", cat_attention.shape)

    data_header_dict['crossAttentions'] = header_cross_attention_list
    data_header_dict['catAttentions'] = header_cat_attention_list
    data_header_dict['labels'] = data_labels
    data_header_dict['tokens'] = header_tokens
    return data_header_dict

def process_sentence_header_data_cellInfo(tokenizer, data_table_id_list, data_header_id_list,
                                 data_highlight_id_list, data_encoder_attentions,
                                 data_cross_attentions):

    print("Process Cell Info!!!!!!")

    data_labels = []
    data_header_dict = {}

    ln_encoder_attentions = []
    for i in data_encoder_attentions:
        i = i.cpu()
        layer_norm_i = nn.LayerNorm([i.shape[1], i.shape[2], i.shape[3]])
        data_encoder_attention = layer_norm_i(i)
        ln_encoder_attentions.append(data_encoder_attention)

    ln_cross_attentions = []
    for i in data_cross_attentions:
        i = i.cpu()
        layer_norm_i = nn.LayerNorm([i.shape[1], i.shape[2], i.shape[3]])
        data_cross_attention = layer_norm_i(i)
        ln_cross_attentions.append(data_cross_attention)

    src_text_id_list = data_table_id_list
    header_text_id_list = data_header_id_list
    highlight_text_id_list = data_highlight_id_list

    # print("src_text_list: ", tokenizer.decode(src_text_id_list))
    # print("len(src_text_list): ", len(tokenizer.decode(src_text_id_list)))
    # print("len(src_id_list): ", len(src_text_id_list))
    # print("src_text_id_list: ", src_text_id_list)
    # print("header_text_list: ", tokenizer.decode(header_text_id_list))
    # print("highlight_text_list: ", tokenizer.decode(highlight_text_id_list,
    #                                                 ))
    # print("highlight_text_id_list", highlight_text_id_list)

    print("check special tokens: ", tokenizer.convert_ids_to_tokens(highlight_text_id_list))

    print("__SEP__: ", tokenizer.convert_tokens_to_ids("__SEP__"))

    # print("len highlight id list: ", len(highlight_text_id_list))
    # print("header id list: ", len(header_text_id_list))

    # rm_special_token_ids = [tokenizer.convert_tokens_to_ids('<pad>'),
    #                         tokenizer.convert_tokens_to_ids('__SEP__'),
    #                         tokenizer.convert_tokens_to_ids('<pad>'),
    #                         tokenizer.convert_tokens_to_ids('__None__'),
    #                         tokenizer.convert_tokens_to_ids('__None__'),
    #                         tokenizer.convert_tokens_to_ids('__EOS__'),
    #                         tokenizer.convert_tokens_to_ids('__PAD__'),
    #                         1437]

    # all headers
    # header_text_id_set = set(header_text_id_list)

    # print(rm_special_token_ids)
    # for special_token in rm_special_token_ids:
    #     if special_token in header_text_id_set:
    #         header_text_id_set.remove(special_token)

    # print("len set highlight id list: ", len(header_text_id_set))
    # get header position in src text
    # position = []
    # for i, src_token_id in enumerate(src_text_id_list):
    #     if src_token_id in header_text_id_set:
    #         position.append(i)
    #
    # print("positions: ", position)
    #
    # eos_position = []
    # eos_id = tokenizer.convert_tokens_to_ids('__EOS__')
    # print(eos_id)
    # for i in range(len(src_text_id_list)):
    #     if src_text_id_list[i] == eos_id:
    #         eos_position.append(i)
    # for i, eos_id in enumerate(src_text_id_list):
    #     if eos_id == src_text_id_list[i]:
    #         eos_position.append(i)

    # print(eos_position)
    #
    # header_cross_attention_list = []
    # header_cat_attention_list = []
    # header_tokens = []
    #
    # start_end_positions = sorted(position+eos_position)
    # print(start_end_positions)
    #
    # print("split header cells")
    # cells_list = []
    # for i in range(len(start_end_positions)):
    #     if i != len(start_end_positions)-1:
    #         if src_text_id_list[start_end_positions[i]] != eos_id:
    #             cell = src_text_id_list[start_end_positions[i]:start_end_positions[i+1]]
    #             cells_list.append(cell)
    #
    #
    #
    #
    # for cell in cells_list:
    #     print(tokenizer.convert_ids_to_tokens(cell))
    #
    #
    # # extract all headers' attentions & labels
    # for i in position:
    #     if (src_text_id_list[i]) in highlight_text_id_list:
    #         data_labels.append(1)
    #     else:
    #         data_labels.append(0)
    #
    #     header_token = tokenizer.convert_ids_to_tokens(src_text_id_list[i])
    #     print(header_token)
    #     header_tokens.append(header_token)

        # encoder_attentions = []
        # for encoder_attention in ln_encoder_attentions:
        #     encoder_attention = encoder_attention.squeeze(0)
        #     header_encoder_attention_i = encoder_attention[:, :, i]
        #     encoder_attentions.append(header_encoder_attention_i)
        #
        # cross_attentions = []
        # for cross_attention in ln_cross_attentions:
        #     cross_attention = cross_attention.squeeze(0)
        #     header_cross_attention_i = cross_attention[:, :, i]
        #     cross_attentions.append(header_cross_attention_i)
        #
        # cat_encoder_attentions = encoder_attentions[0]
        # for encoder_a in encoder_attentions[1:]:
        #     cat_encoder_attentions = torch.cat([cat_encoder_attentions, encoder_a],
        #                                        dim=0)
        #
        # cat_cross_attentions = cross_attentions[0]
        # for cross_a in cross_attentions[1:]:
        #     cat_cross_attentions = torch.cat([cat_cross_attentions, cross_a], dim=0)
        #
        # header_cross_attention_PIL = cvtOneTensorToPIL(cat_cross_attentions)
        # header_cross_attention_list.append(header_cross_attention_PIL)

        # cat_attention = torch.cat((cat_encoder_attentions, cat_cross_attentions),
        #                                                     dim=1)

        # cat_attention_PIL = cvtOneTensorToPIL(cat_attention)
        # header_cat_attention_list.append(cat_attention_PIL)

        # print("cat_encoder_attentions shape ", cat_encoder_attentions.shape)
        # print("cat_cross_attentions shape ", cat_cross_attentions.shape)
        # print("cat_attention shape ", cat_attention.shape)

    # data_header_dict['crossAttentions'] = header_cross_attention_list
    # # data_header_dict['catAttentions'] = header_cat_attention_list
    # data_header_dict['labels'] = data_labels
    # data_header_dict['tokens'] = header_tokens
    # return data_header_dict

def eval_sciGen(text_out_path, reference_path):
    bleu_info = eval_sacre_bleu(reference_path, text_out_path)
    # moverScore = eval_mover_score(reference_path, text_out_path)
    # return bleu_info, moverScore
    return bleu_info
def eval_sacre_bleu(ref_file, pred_file):
    try:
        refs = [get_lines(ref_file)]
        sys = get_lines(pred_file)
        bleu = scb.corpus_bleu(sys, refs)
        return bleu.score
    except:
        return 0

# def eval_mover_score(ref_file, pred_file):
#     try:
#         refs = get_lines(ref_file)
#         sys = get_lines(pred_file)
#         idf_dict_hyp = get_idf_dict(sys)
#         idf_dict_ref = get_idf_dict(refs)
#
#         scores = word_mover_score(refs, sys, idf_dict_ref, idf_dict_hyp, \
#                           stop_words=[], n_gram=1, remove_subwords=True, batch_size=64)
#         return round(np.mean(scores),3) , round(np.median(scores),3 )
#     except Exception as e:
#         print(e)
#         return 0, 0

def get_lines(fil):
    lines = []
    with open(fil, 'r') as f:
        for line in f:
            if line.strip():
                lines.append(line.strip())
            else:
                lines.append('empty')
    return lines


import numpy as np
from pycocotools.coco import COCO
from matplotlib import pyplot as plt
from matplotlib import rcParams
from matplotlib.ticker import ScalarFormatter
from matplotlib import cm
import matplotlib
from collections import defaultdict
from shapely.geometry import Polygon
from matplotlib.font_manager import FontProperties

import wwtool

# plt.rcParams.update({'font.size': 14})    # ICPR paper
plt.rcParams.update({'font.size': 14})

plt.rcParams["font.family"] = "Times New Roman"

# simsun = FontProperties(fname="/usr/share/fonts/msfonts/simsun.ttc")

# matplotlib.rc("font", family='SimSun')
# plt.rcParams["font.family"] = "Microsoft YaHei"

class COCO_Statistic():
    def __init__(self,
                ann_file,
                size_set=[16*16, 32*32, 96*96],
                label_set=[],
                size_measure_by_ratio=False,
                class_instance=None,
                show_title=False,
                out_file_format='pdf',
                max_object_num=2700,
                min_area=0,
                max_small_length=0,
                chinese=False,
                dpi=600):
        self.ann_file = ann_file
        self.coco = COCO(self.ann_file)
        self.catIds = self.coco.getCatIds(catNms=[''])
        self.imgIds = self.coco.getImgIds(catIds=self.catIds)
        self.image_num = len(self.imgIds)
        self.size_set = size_set
        self.label_set = label_set
        self.size_measure_by_ratio = size_measure_by_ratio
        self.class_instance = class_instance
        self.show_title = show_title
        self.out_file_format = out_file_format
        self.max_object_num = max_object_num
        self.min_area = min_area
        self.max_small_length = max_small_length
        self.chinese = chinese
        self.dpi = dpi

        categories = self.coco.dataset['categories']
        self.coco_class = dict()
        for category in categories:
            self.coco_class[category['id']] = category['name']

    def total_size_distribution(self, plot_pie, save_file_name=[]):
        plt.clf()
        plt.figure(figsize=(6, 4.5),)
        object_sizes = []
        image_sizes = []
        size_nums = []
        max_det = 0
        for idx, _ in enumerate(self.imgIds):
            img = self.coco.loadImgs(self.imgIds[idx])[0]
            img_size = img['height'] * img['width']
            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            # print("idx: {}, image file name: {}".format(idx, img['file_name']))
            if len(anns) > max_det:
                max_det = len(anns)
            for ann in anns:
                x1, y1, w, h = ann['bbox']
                if ann['area'] <= self.min_area or max(w, h) < self.max_small_length:
                    continue
                bbox = ann['bbox']
                object_size = bbox[2] * bbox[3]
                if self.size_measure_by_ratio:
                    object_sizes.append(object_size/img_size)
                else:
                    object_sizes.append(object_size)
                image_sizes.append(img_size)
        object_sizes = np.array(object_sizes)
        image_sizes = np.array(image_sizes)

        if plot_pie:
            for idx, size_value in enumerate(self.size_set):
                if idx == 0:
                    size_nums.append(np.where(object_sizes <= size_value)[0].shape[0])
                elif idx != len(self.size_set) - 1:
                    size_nums.append(np.where(object_sizes <= size_value)[0].shape[0] - sum(size_nums[:]))
                else:
                    size_nums.append(len(object_sizes) - sum(size_nums[:]))

            print("size_nums: {}".format(size_nums))

            # plot
            if len(self.label_set) == 0:
                self.label_set = [str(_) for _ in self.size_set]

            fracs = size_nums[:]
            if False:
                explode = [0] * len(self.size_set)
                explode[0] = 0.1
                plt.axes(aspect=1)
                colors = cm.Set1(np.arange(len(self.size_set)) / float(len(self.size_set)))
                plt.pie(x = fracs, labels = self.label_set, explode = explode, colors = colors, autopct = '%3.1f %%', startangle = 0, pctdistance = 0.6)
            else:
                print("size set: ", self.label_set, "size number: ", fracs)
                print("size set: ", self.label_set, "size ratio: ", np.array(fracs)/sum(fracs))
                plt.bar(self.label_set, fracs, color='dodgerblue', alpha=0.75)
                plt.xticks(range(len(self.label_set)), self.label_set, rotation=0)
            if self.show_title:
                plt.title('Size ratio of {} in {} Dataset\n{}'.format(save_file_name[0], save_file_name[1], size_nums))
            save_file_name.append('pie')
        else:
            object_sizes = np.sqrt(np.array(object_sizes))
            image_sizes = np.sqrt(np.array(image_sizes))
            print("Max size: {}, Min size: {}".format(np.max(object_sizes), np.min(object_sizes)))
            print("Absolute Mean size: {}, Var size: {}".format(np.mean(object_sizes), np.std(object_sizes, ddof=1)))
            print("Relative Mean size: {}, Var size: {}".format(np.mean(object_sizes/image_sizes), np.std(object_sizes/image_sizes)))
            # sns_plot = sns.distplot(object_sizes, color="b", bins=30, kde_kws={"lw": 2})
            plt.hist(object_sizes, bins=np.arange(0, 64, 64//30), histtype='bar', facecolor='dodgerblue', alpha=0.75, rwidth=0.95)
            print("Total objects: {}".format(object_sizes.shape[0]))
            print("Max Det: ", max_det)
            if self.show_title:
                plt.title('Size Distribution of {} in\n{} Dataset'.format(save_file_name[0], save_file_name[1]))
            ax = plt.gca()
            xfmt = ScalarFormatter(useMathText=True)
            xfmt.set_powerlimits((0, 0))
            ax.yaxis.set_major_formatter(xfmt)
            if self.chinese:
                plt.xlabel("目标绝对尺度", fontproperties='SimSun')
                plt.ylabel("目标数量", fontproperties='SimSun')
            else:
                plt.xlabel("Instances' sizes")
                plt.ylabel("Instance Count")

            save_file_name.append('hist')

        plt.savefig('{}.{}'.format("_".join(save_file_name), self.out_file_format), bbox_inches='tight', dpi=self.dpi, pad_inches=0.1)
        # plt.show()

    def class_size_distribution(self, coco_class=None, save_file_name=[], number=True):
        plt.clf()
        plt.figure(figsize=(6, 4.5))
        class_nums = defaultdict(lambda: 0)
        class_size = defaultdict(list)
        for idx, _ in enumerate(self.imgIds):
            img = self.coco.loadImgs(self.imgIds[idx])[0]

            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            # print("idx: {}, image file name: {}".format(idx, img['file_name']))
            for ann in anns:
                x1, y1, w, h = ann['bbox']
                if ann['area'] <= self.min_area or max(w, h) < self.max_small_length:
                    continue
                if coco_class == None:
                    coco_class = self.coco_class
                class_name = coco_class[ann['category_id']]
                area = np.sqrt(ann['area'])
                class_size[class_name].append(area)
                class_nums[class_name] += 1

        class_size ={key:class_size[key] for key in sorted(class_size)}
        class_nums ={key:class_nums[key] for key in sorted(class_nums)}

        if number:
            bar_name = np.array(list(class_nums.keys()) if self.class_instance == None else self.class_instance.full2abbr(list(class_nums.keys())))
            bar_value = np.array(list(class_nums.values()))
            sorted_indices = np.argsort(bar_value)[::-1]
            plt.bar(bar_name[sorted_indices], bar_value[sorted_indices], color='dodgerblue', alpha=0.75)
            plt.xticks(range(len(bar_name)), bar_name[sorted_indices], rotation=0)
            if self.chinese:
                plt.xlabel('类别', fontproperties='SimSun')
                plt.ylabel('目标数量', fontproperties='SimSun')
                plt.yscale('log')
            else:
                plt.xlabel('Class')
                plt.ylabel('Instance Count')
                plt.yscale('log')
            save_file_name.append('bar')
        else:
            box_name = list(class_size.keys()) if self.class_instance == None else self.class_instance.full2abbr(list(class_size.keys()))
            box_value = class_size.values()

            bplot = plt.boxplot(box_value, vert=True, whis=1.5, widths=0.6, patch_artist=True, showfliers=False, showmeans=True, labels=box_name, notch=True)
            # plt.xticks(range(1, len(box_name) + 1), box_name, rotation=0)
            if self.chinese:
                plt.ylabel('目标尺度', fontproperties='SimSun')
                plt.xlabel('类别', fontproperties='SimSun')
            else:
                plt.ylabel("Instances' sizes")
                plt.xlabel('Class')
            plt.grid(axis='y')

            # fill with colors
            colors = ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8']
            for patch, color in zip(bplot['boxes'], colors):
                patch.set_facecolor(color)
                patch.set_alpha(0.8)
            save_file_name.append('box')

        print("Classes: {}".format(class_nums))
        if self.show_title:
            plt.title('Class Size Distribution of {} in\n{} Dataset'.format(save_file_name[0], save_file_name[1]))

        plt.savefig('{}.{}'.format("_".join(save_file_name), self.out_file_format), bbox_inches='tight', dpi=self.dpi, pad_inches=0.1)
        # plt.show()

    def image_object_num_distribution(self, save_file_name=[]):
        plt.clf()
        plt.figure(figsize=(6, 4.5))
        object_nums = [0] * self.image_num

        for idx, _ in enumerate(self.imgIds):
            img = self.coco.loadImgs(self.imgIds[idx])[0]
            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            object_num = 0
            for ann in anns:
                x1, y1, w, h = ann['bbox']
                if ann['area'] <= self.min_area or max(w, h) < self.max_small_length:
                    continue
                else:
                    object_num += 1
            # print("idx: {}, image file name: {}".format(idx, img['file_name']))
            object_nums[idx] = object_num

        object_nums = np.array(object_nums)

        # object_sizes = np.sqrt(np.array(object_sizes))
        plt.hist(object_nums, bins=np.arange(0, self.max_object_num, self.max_object_num//80), histtype='bar', facecolor='dodgerblue', alpha=0.75, rwidth=0.9)
        print("Max objects: {}, Image num: {}".format(np.max(object_nums), self.image_num))
        if self.show_title:
            plt.title('Object num Distribution of {} in\n{} Dataset'.format(save_file_name[0], save_file_name[1]))
        plt.xlim([-20, self.max_object_num])
        if self.chinese:
            plt.xlabel('目标数量', fontproperties='SimSun')
            plt.ylabel('图像数量', fontproperties='SimSun')
            plt.yscale('log')
        else:
            plt.xlabel('Instances')
            plt.ylabel('Image Count')
            plt.yscale('log')

        save_file_name.append('object_num')

        plt.savefig('{}.{}'.format("_".join(save_file_name), self.out_file_format), bbox_inches='tight', dpi=self.dpi, pad_inches=0.1)
        # plt.show()

    def object_aspect_ratio_distribution(self, save_file_name=[]):
        plt.clf()
        plt.figure(figsize=(6, 4.5))
        object_aspect_ratios = []
        for idx, _ in enumerate(self.imgIds):
            img = self.coco.loadImgs(self.imgIds[idx])[0]
            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            # print("idx: {}, image file name: {}".format(idx, img['file_name']))
            for ann in anns:
                x1, y1, w, h = ann['bbox']
                if ann['area'] <= self.min_area or max(w, h) < self.max_small_length:
                    continue
                bbox = ann['bbox']
                if bbox[3] == 0:
                    continue
                object_aspect_ratio = bbox[2] / float(bbox[3])
                object_aspect_ratios.append(object_aspect_ratio)
        object_aspect_ratios = np.array(object_aspect_ratios)

        # sns_plot = sns.distplot(object_sizes, color="b", bins=30, kde_kws={"lw": 2})
        plt.hist(object_aspect_ratios, bins=np.arange(0, 3, 3/10), histtype='bar', facecolor='dodgerblue', alpha=0.75, rwidth=0.9)
        if self.show_title:
            plt.title('Aspect Ratio Distribution of {} in\n{} Dataset'.format(save_file_name[0], save_file_name[1]))
        plt.xlabel('size')
        save_file_name.append('aspect_ratio')

        plt.savefig('{}.{}'.format("_".join(save_file_name), self.out_file_format), bbox_inches='tight', dpi=self.dpi, pad_inches=0.1)
        # plt.show()

    def class_num_per_image(self, coco_class=None, save_file_name=[]):
        plt.clf()
        plt.figure(figsize=(6, 4.5))
        class_num_per_image = defaultdict(lambda : 1)

        for idx, _ in enumerate(self.imgIds):
            class_set_per_image = set()
            img = self.coco.loadImgs(self.imgIds[idx])[0]

            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            # print("idx: {}, image file name: {}".format(idx, img['file_name']))
            for ann in anns:
                x1, y1, w, h = ann['bbox']
                if ann['area'] <= self.min_area or max(w, h) < self.max_small_length:
                    continue
                if coco_class == None:
                    coco_class = self.coco_class
                class_name = coco_class[ann['category_id']]
                class_set_per_image.add(class_name)

            class_num_per_image[str(len(class_set_per_image))] +=1
        print(class_num_per_image)
        plt.hist(class_num_per_image.values(), bins=np.arange(0, 8, 8/8), histtype='bar', facecolor='dodgerblue', alpha=0.75, rwidth=0.9)
        plt.xlabel('class')
        plt.ylabel('Instance Count')
        plt.yscale('log')
        save_file_name.append('class_num_per_image')

        if self.show_title:
            plt.title('Class Numbers Distribution of {} in\n{} Dataset'.format(save_file_name[0], save_file_name[1]))

        plt.savefig('{}.{}'.format("_".join(save_file_name), self.out_file_format), bbox_inches='tight', dpi=self.dpi, pad_inches=0.1)
        # plt.show()

    def mask_ratio(self):
        obb_ratio = []
        hbb_ratio = []
        for idx, _ in enumerate(self.imgIds):
            img = self.coco.loadImgs(self.imgIds[idx])[0]
            annIds = self.coco.getAnnIds(imgIds = img['id'], catIds = self.catIds, iscrowd = None)
            anns = self.coco.loadAnns(annIds)       # per image
            for ann in anns:
                bbox = ann['bbox']
                segmentation = ann['segmentation']
                polygon = wwtool.mask2polygon(segmentation[0])
                rbbox = polygon.minimum_rotated_rectangle
                polygon_area = polygon.area
                rbbox_area = rbbox.area

                bbox_area = bbox[2] * bbox[3]

                obb_ratio.append(polygon_area/rbbox_area)
                hbb_ratio.append(polygon_area/bbox_area)

        obb_ratio = np.array(obb_ratio)
        hbb_ratio = np.array(hbb_ratio)

        print("mean obb ratio: {}, mean hbb ratio: {}".format(obb_ratio.mean(), hbb_ratio.mean()))


import torch
from matplotlib import pyplot as plt
from .gru import GRUNet
from torch import nn

def train_gru(data, input_weights, hidden_weights, hidden_size=100,
              start_idx=0, end_idx=1000, verbose=False, reg=0.01, targets=None, washout=0):
    train_data = torch.tensor(data[start_idx:end_idx]).reshape(1, -1, 1).float()
    if targets is None:
        targets = torch.tensor(data[start_idx + 1:end_idx + 1]).reshape(-1, 1).float()[washout:]
    else:
        targets = torch.tensor(targets[start_idx:end_idx]).reshape(-1, 1).float()[washout:]
    model = GRUNet(1, hidden_size, 1, washout=washout)
    model.train()
    # A bit dirty: Set the reservoir weights of the GRU layer to assure they are the same
    model.gru.weight_ih_l0 = nn.Parameter(input_weights.reshape(3 * hidden_size, -1), requires_grad=True)
    model.gru.weight_hh_l0 = nn.Parameter(hidden_weights.reshape(3 * hidden_size, -1), requires_grad=True)
    h = model.init_hidden(1).float()

    hidden_states, h_new = model.gru(train_data, h)
    hidden_states = hidden_states.reshape(-1, hidden_size)[washout:]
    output_weights = torch.linalg.solve(torch.mm(hidden_states.T, hidden_states) + reg * torch.eye(hidden_size),
                                        torch.mm(hidden_states.T, targets))
    model.w_out = output_weights
    prediction_tr, _ = model(train_data, h)
    training_loss = torch.nn.functional.mse_loss(prediction_tr, targets)
    if verbose:
        print("GRU Training complete. Loss:", training_loss.item())
    return model, hidden_states, h_new, training_loss.item()

def validate_gru(y, input_weights, hidden_weights, hidden_size=100,
                 start_idx=0, end_idx=1000, plot=False, targets=None, washout=0):
    train_end_idx = int((end_idx-1)/2)
    test_start_idx = train_end_idx + 1
    model, _, last_hidden, train_loss = train_gru(y, input_weights, hidden_weights, hidden_size,
                                                  start_idx=start_idx, end_idx=train_end_idx, targets=targets)
    test_data = torch.tensor(y[test_start_idx:end_idx-1]).reshape(1, -1, 1).float()
    if targets is None:
      test_targets = torch.tensor(y[test_start_idx+1:end_idx]).float().reshape(-1, 1)[washout:]
    else:
      test_targets = torch.tensor(targets[test_start_idx:end_idx-1]).float().reshape(-1, 1)[washout:]
    model.eval()
    prediction_te, h = model(test_data, last_hidden)
    mse = torch.nn.functional.mse_loss(prediction_te, test_targets)
    #print("MSE: ", mse)
    if plot:
        plt.figure(figsize=(25,4))
        plt.plot(test_targets.detach().numpy(), label="Targets Test")
        plt.plot(prediction_te.detach().numpy(), label="Prediction Test")
        plt.title("GRU")
        plt.legend()
        plt.show()
    return train_loss, mse.item()