import os
import numpy as np

from dataclasses import dataclass
from typing import Dict, Union, NamedTuple

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, RandomSampler, SequentialSampler, DataLoader
from torch.utils.data.distributed import DistributedSampler

from pyutils.display import maybe_tqdm
import pyutils.io as io

from nlpr.tasks.core import BatchMixin
from nlpr.shared.pycore import ExtendedDataClassMixin
from nlpr.shared.model_setup import OptimizerScheduler
from nlpr.shared.modeling.models import forward_batch_delegate, compute_loss_from_model_output
import nlpr.tasks.evaluate as evaluate

class BaseRunner:
    pass

@dataclass
class TrainGlobalState(ExtendedDataClassMixin):
    epoch: int = 0
    epoch_step: int = 0
    global_step: int = 0

    def step(self):
        self.global_step += 1
        self.epoch_step += 1

    def step_epoch(self):
        self.epoch += 1
        self.epoch_step = 0

    def __str__(self):
        return f"TGS({self.epoch} / {self.epoch_step} ({self.global_step}))"

@dataclass
class DatasetWithMetadata:
    dataset: TensorDataset
    metadata: Dict

    def get_descriptor_dict(self):
        return {
            descriptor.name: descriptor
            for descriptor in self.metadata["descriptors"]
        }

@dataclass
class DataDescriptor:
    category: str
    # category=dataset: Tensor
    # category=other: metadata to carry around be batch
    # category=none: don't propagate
    name: str
    pos: Union[int, None]
    # Tensor position. Only use if category=dataset

def convert_examples_to_dataset(examples, tokenizer, feat_spec, task, verbose=False):
    data_rows = [
        example.tokenize(tokenizer).featurize(tokenizer, feat_spec)
        for example in maybe_tqdm(examples, desc="Tokenizing", verbose=verbose)
    ]
    full_batch = task.Batch.from_data_rows(data_rows)
    dataset_with_metadata = full_batch_to_dataset(full_batch)
    dataset_with_metadata.metadata["descriptors"].append(
        DataDescriptor("other_metadata", "example_id", None)
    )
    dataset_with_metadata.metadata["other"]["example_id"] = list(range(len(examples)))
    return dataset_with_metadata

def full_batch_to_dataset(full_batch):
    dataset_ls = []
    others_dict = {}
    descriptors = []
    for i, (k, v) in enumerate(full_batch.asdict().items()):
        if isinstance(v, torch.Tensor):
            descriptors.append(DataDescriptor("dataset", k, len(dataset_ls)))
            dataset_ls.append(v)
        elif v is None:
            descriptors.append(DataDescriptor("none", k, None))
        else:
            descriptors.append(DataDescriptor("other_data", k, None))
            others_dict[k] = v
    dataset_ls.append(torch.arange(len(full_batch)))
    return DatasetWithMetadata(
        dataset=TensorDataset(*dataset_ls),
        metadata={
            "descriptors": descriptors,
            "other": others_dict,
        }
    )

class BatchTuple(NamedTuple):
    batch: BatchMixin
    metadata: dict

    def to(self, device):
        return BatchTuple(
            batch=self.batch.to(device),
            metadata=self.metadata,
        )

class HybridLoader:
    def __init__(self, dataloader, metadata, task):
        self.dataloader = dataloader
        self.metadata = metadata
        self.task = task

    @property
    def dataset_with_metadata(self):
        return DatasetWithMetadata(
            dataset=self.dataloader.dataset,
            metadata=self.metadata,
        )

    def __iter__(self):
        for batch in self.dataloader:
            example_ids = batch[-1]
            batch_dict = {}
            batch_metadata_dict = {"example_id": example_ids}
            for descriptor in self.metadata["descriptors"]:
                if descriptor.category == "dataset":
                    batch_dict[descriptor.name] = batch[descriptor.pos]
                elif descriptor.category == "none":
                    batch_dict[descriptor.name] = None
                elif descriptor.category == "other_data":
                    batch_dict[descriptor.name] = [
                        self.metadata["other"][descriptor.name][i]
                        for i in example_ids
                    ]
                elif descriptor.category == "other_metadata":
                    batch_metadata_dict[descriptor.name] = [
                        self.metadata["other"][descriptor.name][i]
                        for i in example_ids
                    ]
                else:
                    raise KeyError(descriptor.category)
            yield BatchTuple(
                batch=self.task.Batch(**batch_dict),
                metadata=batch_metadata_dict,
            )

    def __len__(self):
        return len(self.dataloader)

def get_sampler(dataset, local_rank, force_sequential=False):
    if force_sequential:
        return SequentialSampler(dataset)
    if local_rank == -1:
        return RandomSampler(dataset)
    else:
        return DistributedSampler(dataset)

def run_val(val_examples, val_dataloader,
            model, task, loss_criterion,
            device, local_rank, verbose):
    if not local_rank == -1:
        return
    model.eval()
    total_eval_loss = 0
    nb_eval_steps, nb_eval_examples = 0, 0
    all_logits = []
    for step, (batch, batch_metadata) in enumerate(
            maybe_tqdm(val_dataloader, desc="Evaluating (Val)", verbose=verbose)):
        batch = batch.to(device)

        with torch.no_grad():
            logits = forward_batch_delegate(
                model=model,
                batch=batch,
                omit_label_ids=True,
                task_type=task.TASK_TYPE,
            )[0]
            tmp_eval_loss = compute_loss_from_model_output(
                logits=logits,
                loss_criterion=loss_criterion,
                batch=batch,
                task_type=task.TASK_TYPE,
            )

        logits = logits.detach().cpu().numpy()
        total_eval_loss += tmp_eval_loss.mean().item()

        nb_eval_examples += len(batch)
        nb_eval_steps += 1
        all_logits.append(logits)
    eval_loss = total_eval_loss / nb_eval_steps
    all_logits = np.concatenate(all_logits, axis=0)

    return {
        "logits": all_logits,
        "loss": eval_loss,
        "metrics": evaluate.compute_task_metrics(task, all_logits, val_examples),
    }

def get_train_dataloader(train_examples, task,
                         tokenizer, feat_spec, local_rank, train_batch_size, verbose=True):
    dataset_with_metadata = convert_examples_to_dataset(
        examples=train_examples,
        feat_spec=feat_spec,
        tokenizer=tokenizer,
        task=task,
        verbose=verbose,
    )
    train_sampler = get_sampler(
        dataset=dataset_with_metadata.dataset,
        local_rank=local_rank,
    )
    train_dataloader = DataLoader(
        dataset=dataset_with_metadata.dataset,
        sampler=train_sampler,
        batch_size=train_batch_size,
    )
    return HybridLoader(
        dataloader=train_dataloader,
        metadata=dataset_with_metadata.metadata,
        task=task,
    )

def get_eval_dataloader(eval_examples, task,
                        tokenizer, feat_spec, eval_batch_size):
    dataset_with_metadata = convert_examples_to_dataset(
        examples=eval_examples,
        feat_spec=feat_spec,
        tokenizer=tokenizer,
        task=task,
    )
    eval_sampler = SequentialSampler(dataset_with_metadata.dataset)
    eval_dataloader = DataLoader(
        dataset=dataset_with_metadata.dataset,
        sampler=eval_sampler,
        batch_size=eval_batch_size,
    )
    return HybridLoader(
        dataloader=eval_dataloader,
        metadata=dataset_with_metadata.metadata,
        task=task,
    )

def complex_backpropagate(loss, optimizer, model,
                          fp16, n_gpu, gradient_accumulation_steps, max_grad_norm):
    if n_gpu > 1:
        loss = loss.mean()  # mean() to average on multi-gpu.
    if gradient_accumulation_steps > 1:
        loss = loss / gradient_accumulation_steps
    if fp16:
        # noinspection PyUnresolvedReferences
        from apex import amp
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_grad_norm)
    else:
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
    return loss

def optim_step_grad_accum(optimizer_scheduler: OptimizerScheduler,
                          train_global_state: TrainGlobalState,
                          gradient_accumulation_steps: int):
    if (train_global_state.epoch_step + 1) % gradient_accumulation_steps == 0:
        optimizer_scheduler.step()
        optimizer_scheduler.optimizer.zero_grad()
        train_global_state.global_step += 1

def save_model_with_metadata(model: nn.Module, metadata: dict, output_dir: str, file_name="model"):
    torch.save(
        model.state_dict(),
        os.path.join(output_dir, f"{file_name}.p")
    )
    io.write_json(
        metadata,
        os.path.join(output_dir, f"{file_name}.metadata.json")
    )

def compare_steps_max_steps(step, max_steps):
    return (
        max_steps is not None
        and max_steps != -1
        and step >= max_steps
    )


"""
Contains functions for building the forest.

Created on Thu May 11 16:30:11 2023

@author: MLechner
# -*- coding: utf-8 -*-
"""
from copy import deepcopy
from numba import njit

import numpy as np
import ray
# import pandas as pd

from mcf import mcf_forest_data_functions as mcf_data
from mcf import mcf_general as mcf_gp
from mcf import mcf_general_sys as mcf_sys
from mcf import mcf_print_stats_functions as ps

def rnd_variable_for_split(x_ind_pos, x_ai_ind_pos, cf_dic, mmm, rng):
    """Generate variables to be used for split.

    Parameters
    ----------
    x_ind_pos : List. Indices of all x-variables.
    x_ai_ind : Numpy 1D array or empty list.
         Indices of all x-variables always used for splitting.
    c_dict : Dict. Parameters
    mmm : Number of variables to draw.
    rng : default random number generator.

    Returns
    -------
    x_i_for_split : List of indices in x of splitting variables.

    """
    qqq = len(x_ind_pos)
    if cf_dic['m_random_poisson'] and mmm > cf_dic['m_random_poisson_min']:
        m_l = 1 + rng.poisson(lam=mmm-1, size=1)
        if m_l < 1:
            m_l = 1
        elif m_l > qqq:
            m_l = qqq
    else:
        m_l = mmm
    if len(x_ai_ind_pos) == 0:
        x_i_for_split = rng.choice(x_ind_pos, m_l, replace=False)
        x_i_for_split_list = x_i_for_split.tolist()
    else:
        if m_l > len(x_ai_ind_pos):
            x_i_for_split = rng.choice(x_ind_pos, m_l-len(x_ai_ind_pos),
                                       replace=False)
            x_i_for_split = np.concatenate((x_i_for_split, x_ai_ind_pos))
            x_i_for_split = np.unique(x_i_for_split)
            x_i_for_split_list = x_i_for_split.tolist()
        else:
            x_i_for_split_list = x_ai_ind_pos[:]
    return x_i_for_split_list

def init_node_table(n_tr, n_oob, indices_oob):
    """Initialise node table for first leaf.

    Parameters
    ----------
    n_tr : INT. Number of observation in training subsample.
    n_oob : INT. Number of observation in OOB subsample.
    indices_oob: Int.

    Returns
    -------
    node_table : List of lists. First init_node_table

    """
    id_node_0 = 0
    id_parent_1 = id_child_left_2 = id_child_right_3 = None
    active_4 = 2
    leaf_size_tr_5, leaf_size_oob_6 = n_tr, n_oob
    objective_fct_value_oob_7 = next_split_i_8 = cut_off_prime_l_9 = None
    x_type_10 = None
    data_tr_indi_11, data_oob_indi_12 = list(range(n_tr)), list(range(n_oob))
    pot_outcomes_13 = pot_variables_used_indi_14 = leaf_size_pot_15 = None
    indices_oob_16 = indices_oob
    node_table = [
        id_node_0, id_parent_1, id_child_left_2, id_child_right_3, active_4,
        leaf_size_tr_5, leaf_size_oob_6, objective_fct_value_oob_7,
        next_split_i_8, cut_off_prime_l_9, x_type_10, data_tr_indi_11,
        data_oob_indi_12, pot_outcomes_13, pot_variables_used_indi_14,
        leaf_size_pot_15, indices_oob_16]
    return [node_table]

def match_cont(d_grid, y_nn, grid_values, rng):
    """
    Select suitable match in case of continuous treatment.

    Parameters
    ----------
    d_grid : Numpy array.
        Discretised treatment.
    y_nn : Numpy array.
        Neighbours.
    leaf_l : Numpy array.
        Observations going to left leaf. (d < larger splitting value).
    leaf_r : Numpy array.
        Observations going to right leaf (d >= larger splitting value).
    grid_values : Numpy array.
        Values (midpoints) used to generate discretised treatment
    rng : Default random number generator object

    Returns
    -------
    y_nn_cont: N x 2 Numpy array.
        Selected neighbours to the left.
    """
    grid = grid_values[1:]    # Controls are not relevant
    min_d_grid, max_d_grid = np.min(d_grid), np.max(d_grid)
    col_no_min = np.argmin(np.abs(grid-min_d_grid))
    col_no_max = np.argmin(np.abs(grid-max_d_grid))
    indices = np.arange(col_no_min, col_no_max + 1)
    nn_indices = rng.choice(indices, size=d_grid.shape[0])
    y_nn_sel = select_one_row_element(y_nn[:, 1:], nn_indices)
    y_nn_red = np.concatenate((y_nn[:, 0].reshape(-1, 1), y_nn_sel), axis=1)
    return y_nn_red

@njit
def select_one_row_element(data, indices):
    """Randomly find one element per row."""
    data_selected = np.empty((data.shape[0], 1))
    for idx, val in enumerate(indices):
        data_selected[idx, 0] = data[idx, val]
    return data_selected

def mcf_mse(y_dat, y_nn, d_dat, w_dat, n_obs, mtot, no_of_treat, treat_values,
            w_yes=False, splitting=False):
    """Compute average mse for the data passed. Based on different methods.

    Parameters
    ----------
    y_dat : Numpy Nx1 vector. Outcome variable of observation.
    y_nn : Numpy N x no_of_treatments array. Matched outcomes.
    d_dat : Numpy Nx1 vector. Treatment.
    w_dat : Numpy Nx1 vector. Weights (or 0)
    n : INT. Leaf size.
    mtot : INT. Method.
    no_of_treat : INT. Number of treated.
    treat_values : List of INT. Treatment values.
    w_yes: Boolean. Weighted estimation.
    splitting: Boolean. Default is False.

    Returns
    -------
    mse : Mean squared error (average not acccount of number of obs).
    treat_share: 1D Numpy array. Treatment shares.

    """
    if w_yes or mtot in (2, 3):
        mse_mce, treat_shares, no_of_obs_by_treat = mcf_mse_not_numba(
            y_dat, y_nn, d_dat, w_dat, n_obs, mtot, no_of_treat,
            treat_values, w_yes, splitting)
    else:
        mse_mce, treat_shares, no_of_obs_by_treat = mcf_mse_numba(
            y_dat, y_nn, d_dat, n_obs, mtot, no_of_treat,
            np.array(treat_values, dtype=np.int8))
    return mse_mce, treat_shares, no_of_obs_by_treat

def mcf_mse_not_numba(y_dat, y_nn, d_dat, w_dat, n_obs, mtot, no_of_treat,
                      treat_values, w_yes, splitting=False):
    """Compute average mse for the data passed. Based on different methods.

    CURRENTLY ONLY USED FOR WEIGHTED.

    Parameters
    ----------
    y_dat : Numpy Nx1 vector. Outcome variable of observation.
    y_nn : Numpy N x no_of_treatments array. Matched outcomes.
    d_all : Numpy Nx1 vector. Treatment.
    w_dat : Numpy Nx1 vector. Weights (or 0)
    n : INT. Leaf size.
    mtot : INT. Method.
    no_of_treat : INT. Number of treated.
    treat_values : List of INT. Treatment values.
    w_yes: Boolean. Weighted estimation.
    splitting: Boolean. Default is False.

    Returns
    -------
    mse : Mean squared error (average not acccount of number of obs).
    treat_share: 1D Numpy array. Treatment shares.

    """
    treat_shares = np.empty(no_of_treat) if mtot in (1, 4) else 0
    mse_mce = np.zeros((no_of_treat, no_of_treat))
    no_of_obs_by_treat = np.zeros(no_of_treat)
    for m_idx in range(no_of_treat):
        d_m = d_dat == treat_values[m_idx]   # d_m is Boolean
        n_m = len(y_dat[d_m])
        no_of_obs_by_treat[m_idx] = n_m
        if w_yes:
            w_m = w_dat[d_m]
            y_m_mean = np.average(y_dat[d_m], weights=w_m, axis=0)
            mse_m = np.average(np.square(y_dat[d_m] - y_m_mean),
                               weights=w_m, axis=0)
        else:
            y_m_mean = np.average(y_dat[d_m], axis=0)
            mse_m = np.dot(y_dat[d_m], y_dat[d_m]) / n_m - (y_m_mean**2)
        if mtot in (1, 4):
            treat_shares[m_idx] = n_m / n_obs
        if mtot in (1, 3, 4):
            mse_mce[m_idx, m_idx] = mse_m
        if mtot != 3:
            mce_ml = 0
            for v_idx in range(m_idx + 1, no_of_treat):
                if mtot == 2:  # Variance of effects mtot = 2
                    d_l = d_dat == treat_values[v_idx]   # d_l is Boolean
                    if w_yes:
                        y_l_mean = np.average(y_dat[d_l], weights=w_dat[d_l],
                                              axis=0)
                    else:
                        y_l_mean = np.average(y_dat[d_l], axis=0)
                    mce_ml = (y_m_mean - y_l_mean)**2
                else:
                    d_ml = (d_dat == treat_values[v_idx]) | (
                        d_dat == treat_values[m_idx])
                    d_ml = d_ml[:, 0]
                    y_nn_m, y_nn_l = y_nn[d_ml, m_idx], y_nn[d_ml, v_idx]
                    if w_yes:
                        w_ml = w_dat[d_ml].reshape(-1)
                        if splitting and (no_of_treat == 2):
                            mce_ml = ((np.average(y_nn_m, weights=w_ml,
                                                  axis=0)) *
                                      (np.average(y_nn_l, weights=w_ml,
                                                  axis=0)) * (-1))
                        else:
                            mce_ml = np.average(
                                (y_nn_m - np.average(y_nn_m, weights=w_ml,
                                                     axis=0)) *
                                (y_nn_l - np.average(y_nn_l, weights=w_ml,
                                                     axis=0)),
                                weights=w_ml, axis=0)
                    else:
                        aaa = np.average(y_nn_m, axis=0) * np.average(y_nn_l,
                                                                      axis=0)
                        bbb = np.dot(y_nn_m, y_nn_l) / len(y_nn_m)
                        mce_ml = bbb - aaa
                mse_mce[m_idx, v_idx] = mce_ml
    return mse_mce, treat_shares, no_of_obs_by_treat

@njit
def mcf_mse_numba(y_dat, y_nn, d_dat, n_obs, mtot, no_of_treat, treat_values):
    """Compute average mse for the data passed. Based on different methods.

       WEIGHTED VERSION DOES NOT YET WORK. TRY with next Numba version.
       Need to change list format soon.

    Parameters
    ----------
    y_dat : Numpy Nx1 vector. Outcome variable of observation.
    y_nn : Numpy N x no_of_treatments array. Matched outcomes.
    d_dat : Numpy Nx1 vector. Treatment.
    d_bin_dat : Numpy Nx1 vector. Treatment larger 0.
    n : INT. Leaf size.
    mtot : INT. Method.
    no_of_treat : INT. Number of treated.
    treat_values : 1D Numpy array of INT. Treatment values.
    cont. Boolean. Continuous treatment.

    Returns
    -------
    mse : Mean squared error (average not acccount of number of obs).
    treat_share: 1D Numpy array. Treatment shares.
    """
    obs = len(y_dat)
    treat_shares = np.zeros(no_of_treat) if mtot in (1, 3, 4) else np.zeros(1)
    mse_mce = np.zeros((no_of_treat, no_of_treat))
    no_of_obs_by_treat = np.zeros(no_of_treat)
    for m_idx in range(no_of_treat):
        d_m = d_dat == treat_values[m_idx]   # d_m is Boolean
        n_m = np.sum(d_m)
        no_of_obs_by_treat[m_idx] = n_m
        y_m = np.empty(n_m)
        j = 0
        for i in range(obs):
            if d_m[i]:
                y_m[j] = y_dat[i, 0]
                j += 1
        y_m_mean = np.sum(y_m) / n_m
        mse_m = np.dot(y_m, y_m) / n_m - (y_m_mean**2)
        if mtot in (1, 3, 4):
            treat_shares[m_idx] = n_m / n_obs
            mse_mce[m_idx, m_idx] = mse_m
        if mtot != 3:
            mce_ml = 0
            for v_idx in range(m_idx + 1, no_of_treat):
                d_l = d_dat == treat_values[v_idx]   # d_l is Boolean
                n_l = np.sum(d_l)
                if mtot == 2:  # Variance of effects mtot = 2
                    y_l = np.empty(n_l)
                    j = 0
                    for i in range(obs):
                        if d_l[i]:
                            y_l[j] = y_dat[i, 0]
                            j += 1
                    y_l_mean = np.sum(y_l) / n_l
                    mce_ml = (y_m_mean - y_l_mean)**2
                elif mtot in (1, 4):
                    d_ml = (d_dat == treat_values[v_idx]) | (
                        d_dat == treat_values[m_idx])
                    n_ml = np.sum(d_ml)
                    y_nn_l = np.empty(n_ml)
                    y_nn_m = np.empty_like(y_nn_l)
                    j = 0
                    for i in range(obs):
                        if d_ml[i]:
                            y_nn_l[j] = y_nn[i, v_idx]
                            y_nn_m[j] = y_nn[i, m_idx]
                            j += 1
                    aaa = np.sum(y_nn_m) / n_ml * np.sum(y_nn_l) / n_ml
                    bbb = np.dot(y_nn_m, y_nn_l) / n_ml
                    mce_ml = bbb - aaa
                mse_mce[m_idx, v_idx] = mce_ml
    return mse_mce, treat_shares, no_of_obs_by_treat

def get_avg_mse_mce(mse_mce, obs_by_treat, mtot, no_of_treat):
    """Bring MSE_MCE matrix in average form."""
    mse_mce_avg = mse_mce.copy()
    for m_idx in range(no_of_treat):
        mse_mce_avg[m_idx, m_idx] = mse_mce[m_idx, m_idx] / obs_by_treat[m_idx]
        if mtot != 3:
            for v_idx in range(m_idx+1, no_of_treat):
                mse_mce_avg[m_idx, v_idx] = mse_mce[m_idx, v_idx] / (
                    obs_by_treat[m_idx] + obs_by_treat[v_idx])
    return mse_mce_avg

def describe_forest(forest, m_n_min_ar, var_dic, cf_dic, gen_dic, pen_mult=0,
                    summary=True):
    """Describe estimated forest by collecting information in trees.

    Parameters
    ----------
    forest : List of List. Each forest consist of one node_table.
    m_n_min : List of INT. Number of variables and minimum leaf size
    v_dict : Dict. Variables.
    c_dict : Dict. Parameters.

    Returns
    -------
    None.

    """
    txt = ('\n' + '-' * 100 + '\nParameters of estimation to build random'
           ' forest')
    txt += '\nOutcome variable used to build forest: '
    txt += ' '.join(var_dic['y_tree_name'])
    txt += '\nFeatures used to build forest:          '
    txt += ' '.join(var_dic['x_name'])
    txt += '\nVariables always included in splitting: '
    txt += ' '.join(var_dic['x_name_always_in'])
    txt += (f'\nNumber of replications:     {cf_dic["boot"]:<4}')
    if cf_dic['mtot'] == 3:
        splitting_rule = 'MSEs of regressions only considered'
    elif cf_dic['mtot'] == 1:
        splitting_rule = 'MSE+MCE criterion'
    elif cf_dic['mtot'] == 2:
        splitting_rule = '-Var(effect)'
    elif cf_dic['mtot'] == 4:
        splitting_rule = 'Random switching'
    txt += f'\nSplitting rule used:        {splitting_rule:<4}'
    if cf_dic['p_diff_penalty'] > 0:
        txt += f'\nPenalty used in splitting:  {pen_mult}'
    txt += '\nShare of data in subsample for forest buildung:'
    txt += f' {cf_dic["subsample_share_forest"]:<4}'
    txt += '\nShare of data in subsample for forest evaluation:'
    txt += f' {cf_dic["subsample_share_eval"]:<4}'
    txt += '\nTotal number of variables available for splitting:'
    txt += f' {len(var_dic["x_name"]):<4}'
    txt += f'\n# of variables (M) used for split: {m_n_min_ar[0]:<4}'
    if cf_dic['m_random_poisson']:
        txt += '\n           (# of variables drawn from 1+Poisson(M-1))'
        txt += '\nMinimum threshold for using Poisson: '
        txt += f'{cf_dic["m_random_poisson_min"]}'
    txt += f'\nMinimum leaf size:                 {m_n_min_ar[1]:<4}'
    txt += f'\nAlpha regularity:                  {m_n_min_ar[2]:5.3f}'
    txt += '\n------------------- Estimated trees ----------------------------'
    leaf_info = get_tree_infos(forest)
    txt += f'\nAverage # of leaves:      {leaf_info[0]:4.1f}'
    txt += f'\nAverage size of leaves:   {leaf_info[1]:4.1f}'
    txt += f'\nMedian size of leaves:    {leaf_info[2]:4.1f}'
    txt += f'\nMin size of leaves:       {leaf_info[3]:4.0f}'
    txt += f'\nMax size of leaves:       {leaf_info[4]:4.0f}'
    txt += f'\nTotal # of obs in leaves: {leaf_info[5]:4.0f}\n' + '-' * 100
    ps.print_mcf(gen_dic, txt, summary=summary)

def get_tree_infos(forest):
    """Obtain some basic information about estimated trees.

    Parameters
    ----------
    forest : List of lists. Collection of node_tables.

    Returns
    -------
    leaf_info : List. Some information about tree.

    """
    leaf_info_tmp = np.zeros([len(forest), 6])
    for boot, tree in enumerate(forest):
        for leaf in tree:
            if leaf[4] == 1:   # Terminal leafs only
                leaf_info_tmp[boot, 0] += 1  # Number of leaves
        leaf_info_tree = np.zeros(int(leaf_info_tmp[boot, 0]))
        j = 0
        for leaf in tree:
            if leaf[4] == 1:
                leaf_info_tree[j] = leaf[5]
                j += 1
        leaf_info_tmp[boot, 1] = np.mean(leaf_info_tree)
        leaf_info_tmp[boot, 2] = np.median(leaf_info_tree)
        leaf_info_tmp[boot, 3] = np.min(leaf_info_tree)
        leaf_info_tmp[boot, 4] = np.max(leaf_info_tree)
        leaf_info_tmp[boot, 5] = np.sum(leaf_info_tree)
    leaf_info = np.empty(6)
    list_of_ind = [0, 1, 5]  # Average #, size of leaves, # of obs in leaves
    leaf_info[list_of_ind] = np.mean(leaf_info_tmp[:, list_of_ind], axis=0)
    leaf_info[2] = np.median(leaf_info_tmp[:, 2])   # Med size of leaves
    leaf_info[3] = np.min(leaf_info_tmp[:, 3])      # Min size of leaves
    leaf_info[4] = np.max(leaf_info_tmp[:, 4])      # Max size of leaves
    return leaf_info

def get_terminal_leaf_no(node_table, x_dat):
    """Get the leaf number of the terminal node for single observation.

    Parameters
    ----------
    node_table : List of list. Single tree.
    x_dat : Numpy array. Data.

    Returns
    -------
    leaf_no : INT. Number of terminal leaf the observation belongs to.

    Note: This only works if nodes are ordered subsequently. Do not remove
          leafs when pruning. Only changes their activity status.

    """
    not_terminal = True
    leaf_id = 0
    while not_terminal:
        leaf = node_table[leaf_id]
        if leaf[4] not in (0, 1):
            raise RuntimeError(f'Leaf is still active. {leaf[4]}')
        if leaf[4] == 1:             # Terminal leaf
            not_terminal = False
            leaf_no = leaf[0]
        elif leaf[4] == 0:          # Intermediate leaf
            if leaf[10] == 0:        # Continuous variable
                leaf_id = (leaf[2] if (x_dat[leaf[8]] - 1e-15) <= leaf[9]
                           else leaf[3])
            else:                   # Categorical variable
                prime_factors = mcf_gp.primes_reverse(leaf[9], False)
                leaf_id = (leaf[2]
                           if int(np.round(x_dat[leaf[8]])) in prime_factors
                           else leaf[3])
    return leaf_no

def remove_oob_from_leaf0(forest):
    """Save memory by removing OOB indices.

    Parameters
    ----------
    forest : List of list. Node_tables.

    Returns
    -------
    forest_out : List of list. Node_tables.
    """
    for idx, _ in enumerate(forest):
        forest[idx][0][16] = 0
    return forest

def fill_trees_with_y_indices_mp(mcf_, data_df, forest):
    """Fill trees with indices of outcomes, MP.

    Returns
    -------
    forest_with_y : List of lists. Updated Node_table.
    terminal_nodes: Tuple of np.arrays. No of final node.
    no_of_avg_nodes: INT. Average no of unfilled leafs.

    """
    int_dic, gen_dic, cf_dic = mcf_.int_dict, mcf_.gen_dict, mcf_.cf_dict
    if int_dic['with_output'] and int_dic['verbose']:
        print("\nFilling trees with indicies of outcomes")
    (x_name, _, _, cf_dic, _, data_np, _, _, x_i, _, _, d_i, _, _, _
     ) = mcf_data.prepare_data_for_forest(mcf_, data_df, True)
    err_txt = 'Wrong order of variables' + str(x_name) + ': ' + str(
        cf_dic['x_name_mcf'])
    if cf_dic['x_name_mcf'] != x_name:
        raise ValueError(err_txt)
    if gen_dic['d_type'] == 'continuous':
        d_dat = data_np[:, d_i]
        # substitute those d used for splitting only that have a zero with
        # random element from the positive treatment levels
        d_pos = d_dat[d_dat > 1e-15]
        rng = np.random.default_rng(12366456)
        d_values = rng.choice(d_pos, size=len(d_dat)-len(d_pos), replace=False)
        d_dat_for_x = np.copy(d_dat)
        j = 0
        for i, d_i in enumerate(d_dat):
            if d_i < 1e-15:
                d_dat_for_x[i, 0] = d_values[j]
                j += 1
        x_dat = np.concatenate((data_np[:, x_i], d_dat_for_x), axis=1)
    else:
        x_dat = data_np[:, x_i]
        d_dat = np.int16(np.round(data_np[:, d_i]))
    obs = len(x_dat)
    terminal_nodes = [None] * cf_dic['boot']
    nodes_empty = np.zeros(cf_dic['boot'])
    if gen_dic['mp_parallel'] < 1.5:
        maxworkers = 1
    else:
        maxworkers = (mcf_sys.find_no_of_workers(gen_dic['mp_parallel'],
                                                 gen_dic['sys_share'])
                      if gen_dic['mp_automatic'] else gen_dic['mp_parallel'])
    if int_dic['with_output'] and int_dic['verbose']:
        print('Number of parallel processes: ', maxworkers)
    if maxworkers == 1:
        for idx in range(cf_dic['boot']):
            (_, forest[idx], terminal_nodes[idx], nodes_empty[idx]
             ) = fill_mp(forest[idx], obs, d_dat, x_dat, idx, gen_dic, cf_dic)
            if int_dic['with_output'] and int_dic['verbose']:
                mcf_gp.share_completed(idx+1, cf_dic['boot'])
    else:
        if int_dic['ray_or_dask'] == 'ray':
            if int_dic['mem_object_store_2'] is None:
                if not ray.is_initialized():
                    ray.init(num_cpus=maxworkers, include_dashboard=False)
            else:
                if not ray.is_initialized():
                    ray.init(num_cpus=maxworkers, include_dashboard=False,
                             object_store_memory=int_dic['mem_object_store_2'])
                if int_dic['with_output'] and int_dic['verbose']:
                    num = round(int_dic["mem_object_store_2"] / (1024 * 1024))
                    txt = f'\nSize of Ray Object Store: {num} MB'
                    ps.print_mcf(gen_dic, txt, summary=False)
            x_dat_ref = ray.put(x_dat)
            still_running = [ray_fill_mp.remote(
                forest[idx], obs, d_dat, x_dat_ref, idx, gen_dic, cf_dic)
                for idx in range(cf_dic['boot'])]
            jdx = 0
            while len(still_running) > 0:
                finished, still_running = ray.wait(still_running)
                finished_res = ray.get(finished)
                for ret_all_i in finished_res:
                    iix = ret_all_i[0]
                    forest[iix] = ret_all_i[1]
                    terminal_nodes[iix] = ret_all_i[2]
                    nodes_empty[iix] = ret_all_i[3]
                    if int_dic['with_output'] and int_dic['verbose']:
                        mcf_gp.share_completed(jdx+1, cf_dic['boot'])
                    jdx += 1
            if 'refs' in int_dic['mp_ray_del']:
                del x_dat_ref
            if 'rest' in int_dic['mp_ray_del']:
                del finished_res, finished
            if int_dic['mp_ray_shutdown']:
                ray.shutdown()
    no_of_avg_enodes = np.mean(nodes_empty)
    if int_dic['with_output'] and int_dic['verbose']:
        txt = ('\nNumber of leaves w/o all treatments per tree: '
               f'{no_of_avg_enodes:6.3%}')
        if no_of_avg_enodes > 0:
            txt += ('\nIncomplete leafs will not be considered for weight'
                    ' computation.')
        txt += '\n' + '-' * 100
        mem = round(mcf_sys.total_size(forest) / (1024 * 1024), 2)
        txt += f'\nSize of forest: {mem} MB' + '\n' + '-' * 100
        ps.print_mcf(gen_dic, txt, summary=True)
    return forest, terminal_nodes, no_of_avg_enodes

@ray.remote
def ray_fill_mp(node_table, obs, d_dat, x_dat, b_idx, gen_dic, cf_dic):
    """Make it work under Ray."""
    return fill_mp(node_table, obs, d_dat, x_dat, b_idx, gen_dic, cf_dic)

def fill_mp(node_table, obs, d_dat, x_dat, b_idx, gen_dic, cf_dic):
    """Compute new node_table and list of final leaves.

    Parameters
    ----------
    node_table : List of lists.
    obs : Int. Sample size.
    d_dat : Numpy array. Treatment.
    x_dat : Numpy array. Features.
    b_idx : Int. Tree number.
    gen_dic, cf_dic : Dict. Controls.

    Returns
    -------
    node_table : List of lists.
    unique_leafs : List.
    b_idx : Int. Tree number.

    """
    subsam = cf_dic['subsample_share_eval'] < 1
    indices = np.arange(obs)
    if subsam:
        obs = round(obs * cf_dic['subsample_share_eval'])
        rng = np.random.default_rng((10+b_idx)**2+121)
        indices = rng.choice(indices, size=obs, replace=False)
    obs_in_leaf = np.zeros((obs, 1), dtype=np.uint32)
    for i, idx in enumerate(indices):
        obs_in_leaf[i] = get_terminal_leaf_no(node_table, x_dat[idx, :])
    unique_leafs = np.unique(obs_in_leaf)
    if subsam:
        unique_leafs = unique_leafs[1:]  # remove first index: obs not used
        d_dat = d_dat[indices]
    nodes_empty = 0
    no_of_treat = (2 if gen_dic['d_type'] == 'continuous'
                   else gen_dic['no_of_treat'])
    for leaf_id in unique_leafs:
        sel_ind = obs_in_leaf.reshape(-1) == leaf_id
        node_table[leaf_id][14] = indices[sel_ind]
        empty_leaf = len(np.unique(d_dat[sel_ind])) < no_of_treat
        if empty_leaf:
            node_table[leaf_id][16] = 1   # Leaf to be ignored
            nodes_empty += 1
    return b_idx, node_table, unique_leafs, nodes_empty/len(unique_leafs)

def save_forests_in_cf_dic(forest_dic, forest_list, fold, no_folds, reg_round,
                           eff_iate):
    """Save forests in dictionary as list of list."""
    # Initialise
    if fold == 0 and reg_round:
        innerlist = [None, None] if eff_iate else [None]
        forest_list = [innerlist for idx in range(no_folds)]
    forest_list[fold][0 if reg_round else 1] = deepcopy(forest_dic)
    return forest_list

def train_save_data(mcf_, data_df, forest):
    """Save data needed in the prediction part of mcf."""
    y_train_df = data_df[mcf_.var_dict['y_name']]
    d_train_df = data_df[mcf_.var_dict['d_name']]
    if mcf_.p_dict['cluster_std']:
        cl_train_df = data_df[mcf_.var_dict['cluster_name']]
    else:
        cl_train_df = None
    if mcf_.gen_dict['weighted']:
        w_train_df = data_df[mcf_.var_dict['w_name']]
    else:
        w_train_df = None
    if mcf_.p_dict['bt_yes']:
        x_bala_train_df = data_df[mcf_.var_dict['x_balance_name']]
    else:
        x_bala_train_df = None
    forest_dic = {'forest': forest, 'y_train_df': y_train_df,
                  'd_train_df': d_train_df, 'x_bala_df': x_bala_train_df,
                  'cl_train_df': cl_train_df, 'w_train_df': w_train_df}
    return forest_dic


"""
Code is partially borrowed from repository https://github.com/sbarratt/inception-score-pytorch/blob/master/inception_score.py # noqa: E501
"""

import argparse
import logging
from pathlib import Path
from typing import Dict, Iterable, Optional, Tuple, Union

import numpy as np
import torch
import torch.utils.data
import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
import yaml
from torch import nn
from torch.nn import functional as F
from torchvision.models.inception import inception_v3
from yaml import Loader

from ex2mcmc.utils.callbacks import Callback, CallbackRegistry

N_INCEPTION_CLASSES = 1000
MEAN_TRASFORM = [0.485, 0.456, 0.406]
STD_TRANSFORM = [0.229, 0.224, 0.225]
N_GEN_IMAGES = 100005000

@CallbackRegistry.register()
class InceptionScoreCallback(Callback):
    def __init__(
        self,
        invoke_every: int = 1,
        device: Union[str, int, torch.device] = "cuda",
        update_input: bool = True,
        dp: bool = False,
        batch_size: int = 128,
    ) -> None:
        self.device = device
        self.model = torchvision.models.inception.inception_v3(
            pretrained=True, transform_input=False
        ).to(device)
        if dp:
            self.model = nn.DataParallel(self.model)

        self.model.eval()
        self.transform = transforms.Normalize(mean=MEAN_TRASFORM, std=STD_TRANSFORM)
        self.update_input = update_input
        self.invoke_every = invoke_every
        self.batch_size = batch_size

    @torch.no_grad()
    def invoke(self, info: Dict[str, Union[float, np.ndarray]]):
        score = None
        step = info.get("step", self.cnt)
        if step % self.invoke_every == 0:
            imgs = torch.from_numpy(info["imgs"]).to(self.device)
            imgs = self.transform(imgs)
            pis = []
            for batch in torch.split(imgs, self.batch_size):
                pis_ = batch_inception(batch, self.model, resize=True).detach().cpu()
                pis.append(pis_)
            pis = torch.cat(pis, 0)
            score = (
                (pis * (torch.log(pis) - torch.log(pis.mean(0)[None, :])))
                .sum(1)
                .mean(0)
            )
            score = torch.exp(score)

            if self.update_input:
                info["inception_score"] = score
            logger = logging.getLogger()
            logger.info(f"\nIS: {score}")
        self.cnt += 1
        return score

def batch_inception(
    imgs: torch.Tensor,
    inception_model: nn.Module,
    resize: bool = False,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    device = imgs.device
    up = nn.Upsample(size=(299, 299), mode="bilinear").to(device)

    def get_pred(x):
        if resize:
            x = up(x)
        x = inception_model(x)
        return F.softmax(x, -1)

    preds = get_pred(imgs)
    return preds

def get_inception_score(
    imgs: Iterable,
    inception_model: Optional[nn.Module] = None,
    gen: Optional[nn.Module] = None,
    generate_from_latents: bool = False,
    cuda: bool = True,
    batch_size: int = 32,
    resize: bool = False,
    splits: int = 1,
    device: Union[torch.device, int] = 0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Computes the inception score of the generated images imgs
    imgs -- Torch dataset of (3xHxW) numpy images normalized in the range [-1, 1]
    cuda -- whether or not to run on GPU
    batch_size -- batch size for feeding into Inception v3
    splits -- number of splits
    """
    N = len(imgs)

    assert batch_size > 0
    assert N >= batch_size

    # Set up dataloader
    if not isinstance(imgs, torch.utils.data.Dataset):
        imgs = torch.utils.data.TensorDataset(imgs)

    dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)

    # Load inception model
    if inception_model is None:
        inception_model = inception_v3(
            pretrained=True,
            transform_input=False,  # False
        ).to(device)
        inception_model.eval()

    up = nn.Upsample(size=(299, 299), mode="bilinear").to(device)

    def get_pred(x):
        if resize:
            x = up(x)
        x = inception_model(x)
        return F.softmax(x, -1).data.cpu()

    # Get predictions
    preds = torch.zeros((N, N_INCEPTION_CLASSES))

    for i, batch in enumerate(dataloader):
        if isinstance(batch, list):
            batch = batch[0]
        batch = batch.to(device)
        if generate_from_latents:
            batch = gen(batch)
        batch_size_i = batch.size()[0]

        preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batch)

    # Now compute the mean kl-div
    split_scores = []

    for k in range(splits):
        part = preds[k * (N // splits) : (k + 1) * (N // splits), :]
        py = torch.mean(part, 0)
        split_scores.append(
            torch.exp(
                (part * (torch.log(part) - torch.log(py[None, :])))
                .sum(1)
                .mean(0)
                # torch.mean(torch.kl_div(part, torch.log(py[None, :])).sum(1))
            )
        )

    return (
        torch.mean(torch.stack(split_scores, 0)),
        torch.std(torch.stack(split_scores, 0)),
        preds,
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--gan_config", type=str)
    parser.add_argument("--device", type=int, default=0)
    args = parser.parse_args()

    device = torch.device(args.device)

    if args.gan_config:
        from ex2mcmc.models.utils import GANWrapper
        from ex2mcmc.utils.general_utils import DotConfig

        gan_config = DotConfig(
            yaml.load(Path(args.gan_config).open("r"), Loader)
        ).gan_config
        # gen, _ = load_gan(gan_config, device)
        gan = GANWrapper(gan_config, device=device, load_weights=True)
        gen = gan.gen

        n_imgs = 1000  # 1000
        imgs = gen(torch.randn(n_imgs, gen.z_dim).to(device))
        imgs = gen.inverse_transform(imgs)

        transform = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        )
        imgs = transform(imgs)

        print("Calculating Inception Score...")
        print(
            get_inception_score(
                imgs,
                cuda=True,
                batch_size=32,
                resize=True,
                splits=1,
            )[:-1]
        )
    else:

        class IgnoreLabelDataset(torch.utils.data.Dataset):
            def __init__(self, orig):
                self.orig = orig

            def __getitem__(self, index):
                return self.orig[index][0]

            def __len__(self):
                return len(self.orig)

        cifar = dset.CIFAR10(
            root="data/",
            download=True,
            transform=transforms.Compose(
                [
                    transforms.Scale(32),
                    transforms.ToTensor(),
                    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                ]
            ),
        )

        IgnoreLabelDataset(cifar)

        print("Calculating Inception Score...")
        print(
            get_inception_score(
                IgnoreLabelDataset(cifar),
                cuda=True,
                batch_size=32,
                resize=True,
                splits=10,
            )[:-1]
        )


import re

from collections import namedtuple
from typing import Any, Callable, Union

from .common import Position, Range

class LexError(Exception):
    pass

Rule = namedtuple('Rule', ['regex', 'callback'])
Token = namedtuple('Token', ['type', 'value', 'range'])

class Lexer:
    """A generic lexer

    Example:
        >>> l = Lexer('aab')
        >>> l.add_rule('a', lambda x: x.upper())
        >>> l.add_rule('b', lambda x: x.upper())
        >>> l.lex()
        'A'
        >>> l.lex()
        'A'
        >>> l.lex()
        'B'
        >>> l.lex()
        None
    """

    def __init__(self, program: str):
        self._program = program
        self._rules = []
        self._lexing = False

    def add_rule(self, regex: str, callback: Callable[[str], Any]):
        """Adds a rule for lexing.

        Args:
            regex: The regular expression string describing the desired lexeme

            callback: A function that processes the matched lexeme text
        """
        if self._lexing:
            raise LexError('Unable to add rules during lexing')

        # Sanitize regex. Make capturing groups non-capturing.
        regex = re.sub('\((?!\?:)', '(?:', regex)

        self._rules.append(Rule(f'({regex})', callback))

    def lex(self) -> Union[Any, None]:
        """Get the next lexeme result.

        Multiple lexemes might be processed until a result is yielded.

        Returns:
            The result of the given callback for the matching lexeme.
            None if the entire source text has been processed.
        """
        try:
            return next(self._generator)

        except AttributeError:
            self._lexing = True

            if not self._rules:
                raise LexError('Cannot lex without rules')

            def generator():
                pattern = re.compile('|'.join([r.regex for r in self._rules]))

                for matches in pattern.findall(self._program):
                    matches = (matches,) if not type(matches) == tuple else matches
                    match, rule = [pair for pair in zip(matches, self._rules) if pair[0]][0]

                    result = rule.callback(match)
                    if result:
                        yield result

                while True:
                    yield None

            self._generator = generator()

        return next(self._generator)

line = 0
character = 0
last_position = Position(line, character)
current_position = Position(line, character)

def update_position(string):
    """Updates the current position of the Lexer.

    Args:
        string: The match string being processed
    """
    global line, character, last_position, current_position
    last_position = Position(line, character)

    for c in string:
        if c == '\n':
            line += 1
            character = 0
        else:
            character += 1

    current_position = Position(line, character)

def process_lexeme(type: str=None) -> Callable[[str], Union[Token, None]]:
    """Construct a function for processing the given type

    Args:
        type: The type of the matched lexeme. If omitted the inner function
            will not return a Token.

    Returns:
        A function
    """

    def callback(lexeme: str) -> Union[Token, None]:
        """Process the given lexeme.

        Args:
            lexeme: The matched lexeme text

        Returns:
            A Token if type was provided by the outer function, otherwise None.
        """
        update_position(lexeme)

        if type:
            range = Range(last_position, current_position)
            return Token(type, lexeme, range)

    return callback

def new_lexer(program: str) -> Lexer:
    """Constructs a Lexer suitable for processing C style structs.

    Args:
        program: The source text

    Returns:
        A Lexer
    """

    global line, character, last_position, current_position
    line = 0
    character = 0
    last_position = Position(line, character)
    current_position = Position(line, character)

    lexer = Lexer(program)

    # Comments
    lexer.add_rule('\/\*([^*]|[\r\n]|(\*+([^*/]|[\r\n])))*\*\/', process_lexeme('comment'))
    #lexer.add_rule('\/\/.*', process_lexeme('comment'))
    lexer.add_rule('\/\/.*(\n\s*\/\/.*)*', process_lexeme('comment'))

    # Literals
    lexer.add_rule('[0-9]+', process_lexeme('number'))

    # Types
    lexer.add_rule('\\b(char|signed char|unsigned char|bool|short|unsigned short|int|unsigned int|long long|unsigned long long|long|unsigned long|float|double)\\b', process_lexeme('type'))

    # Names
    lexer.add_rule('[A-Za-z_]+[A-Za-z0-9_]*', process_lexeme('name'))

    # Operators
    lexer.add_rule('(\[|\]|\{|\}|;|,)', process_lexeme('operator'))

    # Whitespace
    lexer.add_rule('[\s]+', process_lexeme())

    # Everything else. This is for keeping correct positioning.
    lexer.add_rule('.', process_lexeme())

    return lexer


from django.urls import path

from user import views

app_name = 'user'

urlpatterns = [
    path('', views.UserViewSet.as_view(), name='list'),
    path('manage/create/', views.CreateUserView.as_view(), name='create'),
    path('token/', views.CreateTokenView.as_view(), name='token'),
    path('manage/<int:user_id>/', views.ManageUserView.as_view(), name='user'),
]


from datetime import datetime
import pandas as pd
import numpy as np
# import config
import yaml

# time_start = config.time_start
# time_end = config.time_end

with open("config.yaml", 'r') as stream:
    try:
        # print()
        content = yaml.load(stream)
        time_start = content['time_start']
        time_end = content['time_end']

    except yaml.YAMLError as exc:
        print(exc)

def Data_transform():

	df_author = pd.read_csv('group_member_list.csv')

	lan_id_list = df_author['LAN ID'].tolist()

	print(df_author.head())

	print(lan_id_list)
	print(time_start)
	print(time_end)

	df=pd.read_csv('filename.csv')

	df['month-date']=df['timestamp'].apply(lambda x: datetime.utcfromtimestamp(int(x/1000)).strftime('%b-%d'))

	print(df.head())

	df_new = df[(df['year_month_date']>=time_start) & (df['year_month_date']<=time_end) & (df['author_name'].isin(lan_id_list))]

	print(df_new.head())

	df_temp=df_new[['year_month_date','month-date','author_displayName','project','repo']].rename(columns={'author_displayName':'author'})

	df_temp.to_csv('BitBucketDashboardInput.csv',index=False)


#Realsense.py
#Created by: Josh Chapman       Date: 1/12/2023
#This code contains basic implementation for the Realsense camera
#Version 1.0
#Status: Complete
#Capabilities:
    # Camera initialization
    # Getting frames and storing as np arrays
    # ouput of a user-readable image
    # Align depth with color image
    # Select long, medium or short ranges
    # Return a 3D space vector when given a point
    # Generate a point cloud based on a mask like what you would get on numpy.
        #Number of points is limited to ~100
#Abandoned:
    # Calibrate depth image (possibly store to a csv or json file)
        #Skipping this. The camera appears to be well calibrated to start. We can revise this if needed.

#potential fixes:
    #Potentially take only some depth frames to reduce processing load
    #errosion and dialation (remove large noise)
    #Yolo-v7 for AI object detection

#Developed by combining and editing code from:
#https://github.com/IntelRealSense/librealsense/tree/development/wrappers/python/examples
#Techniques for editing detection ranges:
#https://dev.intelrealsense.com/docs/tuning-depth-cameras-for-best-performance
    #This doesn't have any code implementaiton, so will need to find equivalent tools in python library.
    #This page also includes information on error

#Non-Standard Dependencies:
    #Dependency     install command             import line
    #Pyrealsense2   pip install pyrealsense2    import pyrealsense2
    #OpenCV         pip install opencv-python   import cv2

#To get it working on raspbery pi, see:
    # https://dev.intelrealsense.com/docs/using-depth-camera-with-raspberry-pi-3
    # https://www.reddit.com/r/realsense/comments/hn0hfe/anyone_about_to_install_pyrealsense2_on_raspbian/

#Compatibility fix for raspberry pi
try:
    #This block will run on a PC
    import pyrealsense2 as rs
    rs.pipeline() #Call a random RS case to see if the library is imported properly
    pclim = 100 # The maximum number of points that the point cloud will return.
except:
    #The above block will error on a RPi, so we can assume that these variables will be used.
    import pyrealsense2.pyrealsense2 as rs
    pclim = 50 #The maxiumum number of points that the point cloud will return

import cv2
import numpy as np
import time
# from io import StringIO as st
# import pandas as pd

class RSCam:
    '''
    A class for implementing the Intel RealSense D415.

    '''

    def __init__(self,Range = "Long"): #Lower range = lower resolution
        '''
        Lower range = lower resolution. Options are:
        "Long", "Mid", "Short"
        Default is "Long"
        '''
        #Set up RS objects
        self.pc = rs.pointcloud()
        self.pipeline = rs.pipeline()
        self.config = rs.config()

        # Get device product line for setting a supporting resolution
        pipeline_wrapper = rs.pipeline_wrapper(self.pipeline)
        pipeline_profile = self.config.resolve(pipeline_wrapper)
        device = pipeline_profile.get_device()

        # Set the detection range of the camera
        if Range == "Long":
            camX = 1280
            camY = 720
        elif Range == "Mid":
            camX = 640
            camY = 480
        elif Range == "Short":
            camX = 424
            camY = 240

        # Enable the camera
        self.config.enable_stream(rs.stream.depth, camX, camY, rs.format.z16, 30)
        self.config.enable_stream(rs.stream.color, camX, camY, rs.format.bgr8, 30)

        # Start streaming
        self.pipeline.start(self.config)
        self.align = rs.align(rs.stream.color)

    def get_frames(self):
        '''
        Retrieve and align data from the camera.
        Data is stored as RS frames and as numpy arrays for different processing tasks.
        The numpy arrays are returned by this function.
        '''
        while True:
        # Wait for a coherent pair of frames: depth and color
            frames = self.pipeline.wait_for_frames()

            #align frames
            self.frames = self.align.process(frames)
            self.depth_frame = frames.get_depth_frame()
            self.color_frame = frames.get_color_frame()
            if not self.depth_frame or not self.color_frame:
                continue
            else:
                break

        # Convert images to numpy arrays for OpenCV processing
        self.depth_image = np.asanyarray(self.depth_frame.get_data())
        self.color_image = np.asanyarray(self.color_frame.get_data())

        return [self.color_image, self.depth_image]

    def usr_image(self):
        '''
        Create and return a side by side image of the color image and depth map
        '''

        # Change the depth image to a color map to make it easier to see
        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(self.depth_image, alpha=0.03), cv2.COLORMAP_JET)

        # If depth and color resolutions are different, resize color image to match depth image for display
        depth_colormap_dim = depth_colormap.shape
        color_colormap_dim = self.color_image.shape
        if depth_colormap_dim != color_colormap_dim:
            resized_color_image = cv2.resize(self.color_image, dsize=(depth_colormap_dim[1], depth_colormap_dim[0]), interpolation=cv2.INTER_AREA)
            return np.hstack((resized_color_image, depth_colormap)) #Combine Images
        else:
            return np.hstack((self.color_image,depth_colormap)) #Combine Images

    def deproject_Point(self, x = np.zeros(2)):
        '''
        Given a 2D point in the image space, find the 3D point in real space
        # Note: This is a processor intense process. It is better to use FilteredCloud() for multiple points
        '''

        depth_intrin = self.depth_frame.profile.as_video_stream_profile().intrinsics
        depth = self.depth_frame.get_distance(x[0],x[1])
        return rs.rs2_deproject_pixel_to_point(depth_intrin, x, depth)

    def FilteredCloud(self, mask):
        '''
        Return a point cloud of no more than a set number of points specific to the hardware
        being used. The points will corespond to unmasked locations in the input mask.
        '''

        depth = self.depth_frame
        #Generate the point cloud
        self.pc.map_to(self.color_frame)
        points = self.pc.calculate(depth)

        # #Possibly faster method:
        # v, t = points.get_vertices(), points.get_texture_coordinates()
        # verts = np.asanyarray(v).view(np.float32).reshape(-1, 3)  # xyz
        # texcoords = np.asanyarray(t).view(np.float32).reshape(-1, 2)  # uv

        # Bad points appear as 0,0,0 in the point cloud, and correspond to 0 in the depth image.
        # Extracting points is a processor heavy process, because of the way RealSense stores the cloud.
        # Because RealSense isn't officially supported, don't expect this to be fixed.
        # To limit the number of bad points we waste time processing, we will create a mask based on
        # the depth image.

        depthMask = cv2.inRange(self.depth_image,0,1) #This mask finds the bad data points

        #Combine the depth mask with the input mask.
        #Change the arrays to booleans for logical operations
        depthMask = depthMask.astype(bool)
        mask = mask.astype(bool)
        # We want to overlap the good data points with the input mask
        # Since the depth mask is the locations of the bad data points, we just not it
        mask = np.logical_and(mask,np.logical_not(depthMask))

        #Prepare the mask for processing
        mask = mask.flatten() #Make the mask and the point cloud the same dimension
        loc = np.nonzero(mask)#Get the indecies of the non-0 elements of the mask

        #Find the number of points to skip to get an evenly spaced cloud within the allowed number of points
        if np.size(loc) > pclim:
            n = int(np.size(loc)/pclim)
        else:
            n = 1

        #Extract the 3D coordinates
        vtx = np.asarray(points.get_vertices())

        objectCloud = []#initialize empty list
        i = 0 #Initialize the index to 0

        while i < np.size(loc):
            #data type handling:
            line = str(vtx[loc[0][i]])#Convert to a string so we know what the datatype is
            # remove the parenthases from the string
            line = line[1:]
            line = line[:-1]

            temp = line.split(",") #divide the numbers into a list
            #convert the strings to float
            temp = map(float, temp)
            cord = list(temp)

            if sum(cord) != 0: #Check to make sure a bad point didn't slip through
                objectCloud.append(cord) #Good data point, add it and move to the next
                i += n #Skip n coordinates
            else:
                i +=1 #Bad data point, reject it.

        return  np.array(objectCloud)

# # Use area below as example code:
# def nothing(blank):
#     pass

# import matplotlib.pyplot as plt

# cam = RSCam(Range = "Short")
# cv2.namedWindow('RealSense',cv2.WINDOW_AUTOSIZE)

# #assign strings for ease of coding
# hh='Hue High'
# hl='Hue Low'
# sh='Saturation High'
# sl='Saturation Low'
# vh='Value High'
# vl='Value Low'
# wnd = 'Filter Bars'
# cv2.namedWindow(wnd) #Create a window named 'Colorbars'
# #Begin Creating trackbars for each
# cv2.createTrackbar(hl, wnd,0,179,nothing)
# cv2.createTrackbar(hh, wnd,179,179,nothing)
# cv2.createTrackbar(sl, wnd,99,255,nothing)
# cv2.createTrackbar(sh, wnd,255,255,nothing)
# cv2.createTrackbar(vl, wnd,89,255,nothing)
# cv2.createTrackbar(vh, wnd,240,255,nothing)

# #point cloud visualization
# fig = plt.figure(1)
# ax = fig.add_subplot(projection = '3d')

# while(1):
#     cam.get_frames()
#     cam.color_image =  cv2.blur(cam.color_image, [20,20])
#     cv2.imshow('RealSense', cam.usr_image())
#     # print(cam.deproject_Point(np.array([100,100])))

#     #Convert frames to HSV
#     HSV = cv2.cvtColor(cam.color_image, cv2.COLOR_BGR2HSV)

#     #Read trackbars
#     hul=cv2.getTrackbarPos(hl, wnd)
#     huh=cv2.getTrackbarPos(hh, wnd)
#     sal=cv2.getTrackbarPos(sl, wnd)
#     sah=cv2.getTrackbarPos(sh, wnd)
#     val=cv2.getTrackbarPos(vl, wnd)
#     vah=cv2.getTrackbarPos(vh, wnd)

#     #make array for final values
#     HSVLOW=np.array([hul,sal,val])
#     HSVHIGH=np.array([huh,sah,vah])

#     mask1 = cv2.inRange(HSV,HSVLOW, HSVHIGH)
#     #Create a filtered image
#     res = cv2.bitwise_and(cam.color_image,cam.color_image, mask =mask1)

#     cv2.imshow(wnd, res)

#     if cv2.waitKey(1) & 0xFF == ord('b'):
#         break

# while(1):
#     cam.get_frames()
#     HSV = cv2.cvtColor(cam.color_image, cv2.COLOR_BGR2HSV)
#     mask1 = cv2.inRange(HSV,HSVLOW, HSVHIGH)
#     cloud = cam.FilteredCloud(mask1)
#     cv2.imshow('RealSense', cam.usr_image())

#     # plt.scatter(cloud[:][0], cloud[:][1], cloud[:][2])

#     # a = cloud

#     # print(type(cloud[0][0]))
#     if np.size(cloud) > 6:
#         print(np.shape(cloud), "\t", cloud[0:5][1])
#         ax.scatter(cloud[:,0], cloud[:,1], cloud[:,2] )
#         plt.show()
#         cv2.imshow('RealSense', cam.usr_image())
#         break

# while(1):
#     if cv2.waitKey(1) & 0xFF == ord('b'):
#         break


import cv2
import glob
import re
import os
from datetime import datetime

img_array = []
numbers = re.compile(r"(\d+)")
path_to_image_directory = glob.glob('images/ESP32-CAM/*')
latest_image_directory = glob.glob(max(path_to_image_directory, key=os.path.getctime))[0]

def numericalSort(value):
    parts = numbers.split(value)
    parts[1::2] = map(int, parts[1::2])
    return parts

for filename in sorted(glob.glob(latest_image_directory + "/*.jpg") , key=numericalSort):
    img = cv2.imread(filename)
    height, width, layers = img.shape
    size = (width,height)
    img_array.append(img)

out = cv2.VideoWriter("videos/"+ str(datetime.now()) + ".mp4",cv2.VideoWriter_fourcc(*"DIVX"), 15, size)

for i in range(len(img_array)):
    out.write(img_array[i])

out.release()


from django.contrib import admin

# Register your models here.
from blog.models import Profile, Post, Tag

@admin.register(Profile)
class ProfileAdmin(admin.ModelAdmin):
    model = Profile

@admin.register(Tag)
class TagAdmin(admin.ModelAdmin):
    model = Tag

@admin.register(Post)
class PostAdmin(admin.ModelAdmin):
    model = Post

    # In the list of all posts, only show the following information about each post
    list_display = (
        "id",
        "title",
        "subtitle",
        "slug",
        "publish_date",
        "published",
    )
    # Allow filtering the post list by posts that are published or unpublished.
    # Allow filtering posts by publish date.
    list_filter = (
        "published",
        "publish_date",
    )
    # Allow editing all the displayed fields, with the exception of the ID.
    list_editable = (
        "title",
        "subtitle",
        "slug",
        "publish_date",
        "published"
    )
    # Allow searching for posts using the title, subtitle, slug, and body.
    search_fields = (
        "title",
        "subtitle",
        "slug",
        "body"
    )
    # Prepopulate the slug field using the title and subtitle fields.
    prepopulated_fields = {
        "slug": (
            "title",
            "subtitle"
        )
    }
    # Use the publish date of all posts to create a browsable date hierarchy.
    date_hierarchy = "publish_date"
    # Show a button at the top of the list to save changes.
    save_on_top = True


#!/usr/bin/python
# encoding:utf-8

from __future__ import print_function
import json
import sxtwl
import damson
from damson.constraint import (Required, DataType, Between)

from flask import Flask, request

class Result(object):

    @staticmethod
    def fail(code, message):
        return Result(False, code, message)

    @staticmethod
    def sys_error(message=None):
        return Result(False, 'SYS_ERROR', message if message else u'系统繁忙，请稍后再试！')

    @staticmethod
    def params_error(message=None):
        return Result(False, 'PARAMS_ERROR', message if message else u'参数校验未通过！')

    @staticmethod
    def success(data):
        return Result(True, 'SUCCESS', 'ok', data=data)

    def __init__(self, success, code, message, data=None):
        self.success = success
        self.code = code,
        self.message = message
        self.data = data

    def __str__(self):
        result = {
            'success': self.success,
            'message': self.message,
            'code': self.code[0],
            'data': self.data
        }
        return json.dumps(result, ensure_ascii=False)

class OnionLunar:
    TIAN_GAN = ["甲", "乙", "丙", "丁", "戊", "己", "庚", "辛", "壬", "癸"]
    DI_ZHI = ["子", "丑", "寅", "卯", "辰", "巳", "午", "未", "申", "酉", "戌", "亥"]
    SHX = ["鼠", "牛", "虎", "兔", "龙", "蛇", "马", "羊", "猴", "鸡", "狗", "猪"]
    CN_NUM = ["零", "一", "二", "三", "四", "五", "六", "七", "八", "九", "十"]
    JQMC = ["冬至", "小寒", "大寒", "立春", "雨水", "惊蛰", "春分", "清明", "谷雨", "立夏", "小满", "芒种", "夏至", "小暑", "大暑", "立秋", "处暑", "白露",
            "秋分", "寒露", "霜降", "立冬", "小雪", "大雪"]
    YMC = ["冬月", "腊月", "正月", "二月", "三月", "四月", "五月", "六月", "七月", "八月", "九月", "十月"]
    RMC = ["初一", "初二", "初三", "初四", "初五", "初六", "初七", "初八", "初九", "初十", "十一", "十二", "十三", "十四", "十五", "十六", "十七", "十八",
           "十九", "二十", "廿一", "廿二", "廿三", "廿四", "廿五", "廿六", "廿七", "廿八", "廿九", "三十", "卅一"]
    SHI_CHEN = {'子时': ('23:00', '00:59'),
                '丑时': ('01:00', '02:59'),
                '寅时': ('03:59', '04：59'),
                '卯时': ('05:00', '06:59'),
                '辰时': ('07:00', '08:59'),
                '巳时': ('09:00', '10:59'),
                '午时': ('11:00', '12:59'),
                '未时': ('13:00', '14:59'),
                '申时': ('15:00', '16:59'),
                '酉时': ('17:00', '18:59'),
                '戊时': ('19:00', '20:59'),
                '亥时': ('21:00', '22:59')
                }

    def __init__(self):
        self.lunar = sxtwl.Lunar()

    @damson.verify(hour=[DataType(int), Between(0, 24)], minute=[DataType(int), Between(0, 60)])
    def shichen(self, hour, minute):
        key = '%02d:%02d' % (hour, minute)
        for text in self.SHI_CHEN:
            interval = self.SHI_CHEN[text]
            begin = interval[0]
            end = interval[1]
            if begin < end:
                if begin >= key <= end:
                    return text
            else:
                if (begin <= key <= '23:59') or ('00:00' <= key <= end):
                    return text
        return None

    @damson.verify(year=[Required(), DataType(int), Between(1901, 2100, eopen=False)],
                   month=[Required(), DataType(int), Between(1, 12, eopen=False)],
                   day=[Required(), DataType(int), Between(1, 31, eopen=False)])
    def s2y_json(self, year, month, day):
        try:
            day = self.lunar.getDayBySolar(year, month, day)
            data = {
                'year': day.y,
                'gzYear': self.TIAN_GAN[day.Lyear2.tg] + self.DI_ZHI[day.Lyear2.dz],
                'month': day.m,
                'leapMonth': day.Lleap,
                'ymcMonth': self.YMC[day.Lmc],
                'gzMonth': self.TIAN_GAN[day.Lmonth2.tg] + self.DI_ZHI[day.Lmonth2.dz],
                'day': day.d,
                'rmcDay': self.RMC[day.Ldi],
                'gzDay': self.TIAN_GAN[day.Lday2.tg] + self.DI_ZHI[day.Lday2.dz]
            }
            result = Result.success(data)
        except Exception as e:
            result = Result.sys_error(e)
        return str(result)

    @damson.verify(year=[Required(), DataType(int), Between(1901, 2100, eopen=False)],
                   month=[Required(), DataType(int), Between(1, 12, eopen=False)],
                   day=[Required(), DataType(int), Between(1, 31, eopen=False)])
    def y2s_json(self, year, month, day):
        try:
            day = self.lunar.getDayByLunar(year, month, day, True)
            data = {
                'year': day.y,
                'gzYear': self.TIAN_GAN[day.Lyear2.tg] + self.DI_ZHI[day.Lyear2.dz],
                'month': day.m,
                'leapMonth': day.Lleap,
                'ymcMonth': self.YMC[day.Lmc],
                'gzMonth': self.TIAN_GAN[day.Lmonth2.tg] + self.DI_ZHI[day.Lmonth2.dz],
                'day': day.d,
                'rmcDay': self.RMC[day.Ldi],
                'gzDay': self.TIAN_GAN[day.Lday2.tg] + self.DI_ZHI[day.Lday2.dz]
            }
            result = Result.success(data)
        except Exception as e:
            result = Result.sys_error(e)
        return result

# +++++++++++++++++ Flask Web ++++++++++++++++++++

app = Flask(__name__)
onion = OnionLunar()

@app.route('/')
def home():
    html = """
    ©Copyright since 2019 xasync.com

Welcome to use Onion-Lunar!

—— Power by xasync.com

    """
    return html

@app.route('/solar2lunar', methods=['GET'])
def solar2lunar():
    try:
        year = int(request.args.get('year'))
        month = int(request.args.get('month'))
        day = int(request.args.get('day'))
        result = onion.s2y_json(year, month, day)
    except Exception as e:
        result = Result.sys_error(e)
    return str(result)

@app.route('/lunar2solar', methods=['GET'])
def lunar2solar():
    try:
        year = int(request.args.get('year'))
        month = int(request.args.get('month'))
        day = int(request.args.get('day'))
        result = onion.y2s_json(year, month, day)
    except Exception as e:
        result = Result.sys_error(e)
    return str(result)

if __name__ == '__main__':
    app.run('127.0.0.1', port=15954, debug=False)


from __future__ import print_function
from __future__ import division

import logging
logging.basicConfig(format='%(asctime)s | %(levelname)s : %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)

logger.info("Loading packages ...")

import os
import sys
import torch
import numpy as np
from tqdm import tqdm
import json
from datetime import datetime
import time
import traceback
import argparse
import warnings

# Project modules
import utils
import model
import training
import dataset

# Handle command line arguments
parser = argparse.ArgumentParser(description='Run a complete training pipeline. Optionally, a JSON configuration file can be used, to overwrite command-line arguments.')
# I/O
parser.add_argument('--config', dest='config_filepath',
                    help='Configuration .json file (optional). Overwrites existing command-line args!')
parser.add_argument('--output_dir', default='/users/gzerveas/data/gzerveas/pruned_med_img_seg/output',
                    help='Root output directory. Must exist. Time-stamped directories will be created inside.')
parser.add_argument('--data_dir', default='/gpfs/data/ceickhof/GAILA',
                    help='Data directory')
parser.add_argument('--load_model',
                    help='Path to pretrained model')
parser.add_argument('--debug_size', type=int,
                    help='For rapid testing purposes (e.g. debugging), limit training set to a small random sample')
parser.add_argument('--name', dest='experiment_name', default='',
                    help='A string identifier/name for the experiment to be run - it will be appended to the output directory name, before the timestamp')
parser.add_argument('--no_timestamp', action='store_true',
                    help='If set, a timestamp will not be appended to the output directory name')
parser.add_argument('--records_file', default='/gpfs/data/ceickhof/GAILA/records.xls',
                    help='Excel file keeping all records of experiments')
# Training
parser.add_argument('--gpu', type=str, default='0',
                    help='GPU index, -1 for CPU')
parser.add_argument('--num_workers', type=int, default=4,
                    help='dataloader threads. 0 for single-thread.')
parser.add_argument('--epochs', type=int, default=100,
                    help='Number of training epochs')
parser.add_argument('--lr', type=float, default=1.25e-4,
                             help='learning rate for batch size 32.')
parser.add_argument('--lr_step', type=str, default='90,120',
                             help='Comma separated string of epochs when to reduce learning rate by a factor of 10')
parser.add_argument('--batch_size', type=int, default=10,
                    help='Training batch size')
parser.add_argument('--img_size', type=int, default=128,
                    help='Final size of training images (after downsampling)')
parser.add_argument('--seed',
                    help='Seed used for splitting sets. None by default, set to an integer for reproducibility')

args = parser.parse_args()

NUM_CLASSES = 16

if args.task == '3d':
    heads = {'hm': NUM_CLASSES, 'dep': 1, 'rot': 8, 'dim': 3}
    if args.reg_bbox:
        heads.update({'wh': 2})
    if args.reg_offset:
        heads.update({'reg': 2})
elif args.task == 'localization':
    heads = {'hm': NUM_CLASSES, 'wh': 2}

def setup(args):
    """Prepare training session: read configuration from file (takes precedence), create directories.
    Input:
        args: arguments object from argparse
    Returns:
        config: configuration dictionary
    """

    config = args.__dict__  # configuration dictionary

    if args.config_filepath is not None:
        logger.info("Reading configuration ...")
        try:  # dictionary containing the entire configuration settings in a hierarchical fashion
            config.update(utils.load_config(args.config_filepath))
        except:
            logger.critical("Failed to load configuration file. Check JSON syntax and verify that files exist")
            traceback.print_exc()
            sys.exit(1)

    # Create output directory
    initial_timestamp = datetime.now()
    output_dir = config['output_dir']
    if not os.path.isdir(output_dir):
        raise IOError(
            "Root directory '{}', where the directory of the experiment will be created, must exist".format(output_dir))

    output_dir = os.path.join(output_dir, config['experiment_name'])

    formatted_timestamp = initial_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
    config['initial_timestamp'] = formatted_timestamp
    if (not config['no_timestamp']) or (len(config['experiment_name']) == 0):
        output_dir += "_" + formatted_timestamp
    utils.create_dirs([output_dir])
    config['output_dir'] = output_dir

    # Save configuration as a (pretty) json file
    with open(os.path.join(output_dir, 'configuration.json'), 'w') as fp:
        json.dump(config, fp, indent=4, sort_keys=True)

    logger.info("Stored configuration file in '{}'".format(output_dir))

    return config

def main():

    start_time = time.time()

    config = setup(args)

    torch.manual_seed(config['seed'])

    device = torch.device('cuda' if (torch.cuda.is_available() and config['gpu'] != '-1') else 'cpu')

    # Create or load model
    # Fail early if something is amiss
    logger.info("Creating model ...")
    my_model = model.DCN_Detector(opt.heads, opt.head_conv)
    optimizer = torch.optim.Adam(my_model.parameters(), opt.lr)
    start_epoch = 0
    if args.load_model:
        my_model, optimizer, start_epoch = utils.load_model(my_model, args.load_model, optimizer, args.resume, args.lr, args.lr_step)
    my_model.to(device)

    # Initialize data generators
    logger.info("Loading and preprocessing data ...")
    my_dataset = dataset.ObjDetDataset()

    val_loader = torch.utils.data.DataLoader(
        my_dataset(opt, 'val'),
        batch_size=1,
        shuffle=False,
        num_workers=1,
        pin_memory=True
    )

    trainer = training.Trainer(opt, my_model, optimizer)
    trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)

    if opt.test: # Only evaluate and skip training
        _, preds = trainer.val(0, val_loader)
        val_loader.dataset.run_eval(preds, opt.save_dir)
        return

    train_loader = torch.utils.data.DataLoader(
        my_dataset(opt, 'train'),
        batch_size=opt.batch_size,
        shuffle=True,
        num_workers=opt.num_workers,
        pin_memory=True,
        drop_last=True
    )

    logger.info('Starting training...')
    best = 1e10
    metrics = []  # list of lists: for each epoch, stores metrics like accuracy, DICE,
    it = tqdm(range(1, config["epochs"] + 1), desc='Training', ncols=0)
    for epoch in range(start_epoch + 1, opt.num_epochs + 1):
        mark = epoch if opt.save_all else 'last'
        log_dict_train, _ = trainer.train(epoch, train_loader)
        logger.info('epoch: {} |'.format(epoch))
        for k, v in log_dict_train.items():
            logger.info('{} {:8f} | '.format(k, v))
        if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
            save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)),
                        epoch, my_model, optimizer)
            with torch.no_grad():
                log_dict_val, preds = trainer.val(epoch, val_loader)
            for k, v in log_dict_val.items():
                logger.scalar_summary('val_{}'.format(k), v, epoch)
                logger.write('{} {:8f} | '.format(k, v))
            if log_dict_val[opt.metric] < best:
                best = log_dict_val[opt.metric]
                save_model(os.path.join(opt.save_dir, 'model_best.pth'),
                        epoch, my_model)
        else:
            save_model(os.path.join(opt.save_dir, 'model_last.pth'),
                        epoch, my_model, optimizer)
        if epoch in opt.lr_step:
            save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), epoch, my_model, optimizer)
            lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
            print('Drop LR to', lr)
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr

        metrics.append([epoch, accuracy, dice, precision, recall])

        it.set_postfix(
            Accuracy='{:.3f}%'.format(accuracy * 100),
            Max_Accuracy='{:.3f}%'.format(max_acc * 100),
            Dice='{:.3f}%'.format(dice * 100),
            Max_Dice='{:.3f}%'.format(max_dice * 100),
            Precision='{:.3f}%'.format(precision * 100),
            Max_Precision='{:.3f}%'.format(max_prec * 100),
            Recall='{:.3f}%'.format(recall * 100),
            Max_Recall='{:.3f}%'.format(max_rec * 100),
            Last_pruning='{:2d}'.format(latest_pruning),
            Pruned='{:.2f}%'.format(pruned_ratio * 100))

    # Export evolution of metrics over epoch
    header = ["Epoch", "Accuracy", "DICE", "Precision", "Recall", "Pruned ratio"]
    metrics_filepath = os.path.join(config["output_dir"], "metrics_" + config["experiment_name"] + ".xls")
    book = utils.export_performance_metrics(metrics_filepath, metrics, header, sheet_name="metrics")

    # Export best metrics per pruning level
    header = ["Pruned ratio", "DICE", "Accuracy", "Precision", "Recall"]
    utils.export_performance_metrics(metrics_filepath, best_at_prune_level, header, book, sheet_name="best_metrics")

    # Export record metrics to a file accumulating records from all experiments
    utils.register_record(config["records_file"], config["initial_timestamp"], config["experiment_name"], metrics)

    logger.info('All Done!')

    total_runtime = time.time() - start_time
    logger.info(
        "Total runtime: {} hours, {} minutes, {} seconds\n".format(total_runtime // 3600, (total_runtime // 60) % 60,
                                                                   total_runtime % 60))

if __name__ == '__main__':

    main()