import os
#from this import d

import numpy as np
import pandas as pd
from scipy.optimize import minimize, basinhopping

import matplotlib.pyplot as pl
from matplotlib import rcParams

from .. import logging as logg
from .. import settings
from ..core import get_n_jobs, parallelize, SplicingDynamics
from ..preprocessing.moments import get_connectivities
from .dynamical_model_utils import BaseDynamics, convolve, linreg, get_norm_std, use_model
from .utils import make_unique_list, test_bimodality

MAX_LOSS = 1e6

class DynamicsRecovery(BaseDynamics):
    def __init__(self, adata, gene, load_pars=None, **kwargs):
        super().__init__(adata, gene, **kwargs)
        if load_pars and "fit_alpha" in adata.var.keys():
            try:
                self.load_pars(adata, gene)
            except Exception:
                self.recoverable = False
                logg.warn(f"Model for {self.gene} could not be instantiated in load_pars.")

        elif self.recoverable:
            self.initialize()

    def initialize_WX(self, method):
        dim = self.n_TFs
        #tmp = self.gamma - np.sqrt(self.gamma*self.gamma+8*self.beta)
        #max_WX = (self.gamma-tmp/2)*self.alpha*np.exp(-tmp*tmp/(16*self.beta)) + self.gamma*self.rho
        #max_x = self.x_f.sum(1).max()
        #norm_term = max_WX/max_x
        if method == 'rand':
            weight = np.random.rand(dim)*2/dim  #* (2*norm_term)
        elif method == 'ones':
            #weight = np.ones(dim) * norm_term
            #weight = np.ones(dim)/dim
            #weight = (self.y_f.mean()/self.x_f.mean(0))/dim
            weight = np.ones(dim)/self.x_f.sum(1).std()
        elif method == 'knockTF_Log2FC':
            weight = self.knockTF_Log2FC / np.dot(self.x_f, self.knockTF_Log2FC.reshape(-1,1)).std()
        elif method == 'correlation':
            weight = self.TFs_correlation / np.dot(self.x_f, self.TFs_correlation.reshape(-1,1)).std()
        elif method == 'PCA':
            from sklearn.decomposition import PCA
            pca = PCA(n_components=1)
            pca.fit_transform(self.x_f)
            weight = pca.components_[0]
        elif method == 'gamma':
            from scipy.optimize import lsq_linear
            tmp1 = self.max_yf_x.copy()
            tmp2 = self.gamma * self.max_yf
            lb = -self.WX_thres
            hb = self.WX_thres
            weight = lsq_linear(tmp1, tmp2, bounds=(lb, hb), method='bvls', lsmr_tol='auto', verbose=0)
            weight = weight.x
        weight = weight.reshape(-1,1)
        WX = np.ravel(np.dot(self.x, weight))
        return WX, weight

    def init_par(self, alpha=None, beta=None, omega=None, theta=None, gamma=None):
        self.alpha = float(self.max_yf.mean()-self.min_yf.mean())/2 if alpha is None else alpha
        self.beta = float(self.max_yf.mean()+self.min_yf.mean())/2 if beta is None else beta
        self.omega = 2*np.pi if omega is None else omega
        self.theta = 0 if theta is None else theta
        self.gamma = 1 if gamma is None else gamma

        if not self.init_weight_method is None:
            self.WX, self.weights_init = self.initialize_WX(method=self.init_weight_method)
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)

        # update object with initialized vars
        self.likelihood, self.varx = 0, 0

        # initialize time point assignment
        if not self.init_weight_method is None:
            self.WX_model, self.delta = self.update_WX(self.alpha, self.beta, self.omega, self.theta,
                self.gamma, method=self.WX_method, thres=self.WX_thres)
        self.WX = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta)
        max_yf_WX = use_model(self.WX_model, self.max_yf_x, method=self.WX_method, bias=self.delta)
        self.gamma = float(np.mean(max_yf_WX/self.max_yf))
        #if self.gamma < self.gamma_thres:
        #    self.gamma = self.gamma_thres
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)
        self.pars = np.array([self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta,
            self.likelihood, self.varx, self.scaling_y] + list(self.scaling))[:, None]

        self.loss = [self.get_loss(refit_time=False, norm_std=True, reg=self.reg)]
        return

    def initialize(self):
        TFs_x, target_y, f = self.TFs_x, self.target_y, self.filtered
        TFs_x_f = TFs_x[f]
        target_y_f = target_y[f]

        # initialize scaling
        #for x_i in self.TFs_expression
        self.std_TFs_x, self.std_target_y = np.std(TFs_x_f, 0), np.std(target_y_f)
        #self.std_TFs_x[self.std_TFs_x==0] = 1e6 #TFs_x_f[:,self.std_TFs_x==0].mean(0)
        if self.std_target_y == 0:
            self.std_TFs_x = self.std_target_y = 1
        _scaling = self.fit_scaling
        if isinstance(_scaling, bool):
            if _scaling:
                scaling = 1/self.std_TFs_x
                scaling = np.nan_to_num(scaling, nan=0, posinf=0, neginf=0)
                scaling_y = 1/self.std_target_y
            else:
                scaling = np.ones_like(self.std_TFs_x)
                scaling_y = 1
        else:
            scaling = _scaling

        x, x_f = TFs_x * scaling, TFs_x_f * scaling
        y, y_f = target_y * scaling_y, target_y_f * scaling_y
        self.x, self.x_f, self.y, self.y_f = x, x_f, y, y_f
        self.scaling, self.scaling_y = scaling, scaling_y

        self.max_yf_idx = np.ravel(y_f >= np.percentile(y_f, 99, axis=0))
        self.max_yf_x = x_f[self.max_yf_idx]
        self.max_yf = np.ravel(y_f[self.max_yf_idx])
        self.min_yf_idx = np.ravel(y_f <= np.percentile(y_f, 1, axis=0))
        self.min_yf = np.ravel(y_f[self.min_yf_idx])

        self.init_par()
        return

    def fit(self):
        if self.fix_theta:
            self.search_and_fit_no_theta(reg=self.reg)
        else:
            self.search_and_fit(reg=self.reg)
        self.WX = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta)
        self.t = self.assign_time(self.WX, self.y, self.alpha, self.beta, self.omega, self.theta,
            self.gamma, self.filtered, norm_std=True)
        self.varx = self.get_variance(to_filter=True, refit_time=False, norm_std=True)
        self.likelihood = self.get_likelihood(to_filter=True, varx=self.varx, refit_time=False, norm_std=True)
        self.weights = np.ravel(self.WX_model)
        self.weights_final = self.weights * self.x_f.sum(0)
        self.weights_final = self.weights_final / abs(self.weights_final).max()
        #if :
        #    self.WX_model = - self.WX_model
        #    self.fit_paras_(reg=self.reg)
        self.get_raw_velocity()

    def save_par(self):
        self.best_loss = self.loss[-1]
        self.alpha_best, self.beta_best, self.omega_best, self.theta_best, self.gamma_best, self.delta_best = \
            self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta
        self.loss_best = self.loss
        self.t_best = self.t
        self.WX_model_best = self.WX_model
        self.WX_best = self.WX
        return
    def load_par(self):
        self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta =\
            self.alpha_best, self.beta_best, self.omega_best, self.theta_best, self.gamma_best, self.delta_best
        self.loss = self.loss_best
        self.t = self.t_best
        self.WX_model = self.WX_model_best
        self.WX = self.WX_best
        return
    def search_and_fit(self, **kwargs):
        omega_vales = [2*np.pi] #np.linspace(1/4*np.pi, 2*np.pi, num=8)
        self.save_par()
        self.best_loss = MAX_LOSS
        for omega in omega_vales:
            theta_vals = np.linspace(-0.5, 0.5, num=1+int(omega/(np.pi/4))) * omega
            for theta in theta_vals:
                self.init_par(omega=omega, theta=theta)
                self.fit_paras_(**kwargs)
                if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                    continue
                if self.loss[-1] < self.best_loss:
                    self.save_par()
        self.load_par()
        self.tmp_best_loss = self.best_loss

        gamma_vals = np.linspace(2.5, 7.5, num=3)
        for omega in omega_vales:
            theta_vals = np.linspace(-0.5, 0.5, num=1+int(omega/(np.pi/4))) * omega
            for theta in theta_vals:
                for gamma in gamma_vals:
                    if gamma < self.gamma_thres:
                        continue
                    self.init_par(omega=omega, theta=theta, gamma=gamma)
                    self.fit_paras_(**kwargs)
                    if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                        continue
                    if (self.loss[-1] < self.best_loss) and (self.loss[-1] < 0.9*self.tmp_best_loss):
                        self.save_par()
        self.load_par()
        return

    def fit_paras_(self, **kwargs):
        def mse(data):
            #return self.get_mse(alpha=data[0], beta=data[1], omega=data[2], theta=data[3], gamma=data[4], **kwargs)
            return self.get_mse(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], **kwargs)
        #data_0, cb = np.array([self.alpha, self.beta, self.omega, self.theta, self.gamma]), self.cb_fit_paras_
        data_0, cb = np.array([self.alpha, self.beta, self.theta, self.gamma]), self.cb_fit_paras_
        #bounds=((0, None), (None, None), (0, 2*np.pi), (-np.pi, np.pi), (0, None))
        bounds=((0, None), (None, None), (-np.pi, np.pi), (self.gamma_thres, None))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_fit_paras_(res.x)
    def cb_fit_paras_(self, data):
        #self.update(alpha=data[0], beta=data[1], omega=data[2], theta=data[3], gamma=data[4], refit_time=True, update_weight=True)
        self.update(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], refit_time=True, update_weight=True)

    def search_and_fit_no_theta(self, **kwargs):
        omega = 2*np.pi
        theta = 0
        self.save_par()
        self.best_loss = MAX_LOSS
        self.init_par(omega=omega, theta=theta)
        self.fit_paras_no_theta_(**kwargs)
        if self.loss[-1] < self.best_loss:
            self.save_par()
        self.load_par()
        self.tmp_best_loss = self.best_loss

        gamma_vals = np.linspace(2.5, 7.5, num=3)
        for gamma in gamma_vals:
            if gamma < self.gamma_thres:
                continue
            self.init_par(omega=omega, theta=theta, gamma=gamma)
            self.fit_paras_no_theta_(**kwargs)
            if (self.best_loss<MAX_LOSS) and ((self.WX.max() - self.WX.min()) < 1/3 * (2*self.alpha*np.sqrt(4*np.pi**2 + self.gamma**2))):
                continue
            if (self.loss[-1] < self.best_loss) and (self.loss[-1] < 0.9*self.tmp_best_loss):
                self.save_par()
        self.load_par()
        return

    def fit_paras_no_theta_(self, **kwargs):
        def mse(data):
            return self.get_mse(alpha=data[0], beta=data[1], gamma=data[2], **kwargs)
        data_0, cb = np.array([self.alpha, self.beta, self.gamma]), self.cb_fit_paras_no_theta_
        bounds=((0, None), (None, None), (self.gamma_thres, None))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_fit_paras_no_theta_(res.x)
    def cb_fit_paras_no_theta_(self, data):
        self.update(alpha=data[0], beta=data[1], gamma=data[2], refit_time=True, update_weight=True)

    def refit_paras_(self, **kwargs):
        def mse(data):
            return self.get_mse(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], delta=data[4], **kwargs)
        data_0, cb = np.array([self.alpha, self.beta, self.theta, self.gamma, self.delta]), self.cb_refit_paras_
        bounds=((0, None), (None, None), (-np.pi, np.pi), (0, None), (-10, 10))
        res = minimize(mse, data_0, callback=cb, bounds=bounds, **self.simplex_kwargs)
        self.cb_refit_paras_(res.x)
    def cb_refit_paras_(self, data):
        self.update(alpha=data[0], beta=data[1], theta=data[2], gamma=data[3], delta=data[4], refit_time=True, update_weight=False)

    def update(
        self,
        t=None,
        alpha=None,
        beta=None,
        omega=None,
        theta=None,
        gamma=None,
        delta=None,
        refit_time=False,
        update_weight=True,
    ):
        loss_prev = self.loss[-1] if len(self.loss) > 0 else MAX_LOSS

        alpha, beta, omega, theta, gamma, delta = self.get_vars(alpha, beta, omega, theta, gamma, delta)

        if update_weight:
            WX_model, delta = self.update_WX(alpha, beta, omega, theta, gamma, method=self.WX_method, thres=self.WX_thres)
        else:
            WX_model = self.WX_model

        WX = use_model(WX_model, self.x, method=self.WX_method, bias=delta)
        if refit_time:
            t = self.assign_time(WX, self.y, alpha, beta, omega, theta, gamma, self.filtered, norm_std=True)
        else:
            t = self.t

        loss = self.get_loss(t, alpha, beta, omega, theta, gamma, delta, WX_model, refit_time=False, norm_std=True, reg=self.reg)
        perform_update = loss < loss_prev

        if perform_update:
            self.alpha, self.beta, self.omega, self.theta, self.gamma, self.delta = alpha, beta, omega, theta, gamma, delta
            self.t = t
            self.WX = WX
            self.WX_model = WX_model
            new_pars = np.array([alpha, beta, omega, theta, gamma, delta, self.likelihood, self.varx, self.scaling_y] + self.scaling.reshape(-1).tolist())[:, None]
            self.pars = np.c_[self.pars, new_pars]
            self.loss.append(loss)
            #print('loss:',loss)

        return perform_update

    def assign_time(
        self, WX, y, alpha, beta, omega, theta, gamma, filtered, norm_std=False
    ):
        num = np.clip(len(y), self.n_time_points, self.n_time_points*2)
        tpoints = np.linspace(0, 1, num=num)
        data_model = SplicingDynamics(alpha=alpha, beta=beta, omega=omega, theta=theta,
            gamma=gamma).get_solution(tpoints)

        if norm_std:
            std_WX = get_norm_std(data_model[:,0], WX[filtered]) # WX[filtered], data_model[:,0]
            std_y = get_norm_std(data_model[:,1], y[filtered], par=1.5) # y[filtered], data_model[:,1]
        else:
            std_WX, std_y = 1, 1
        data_model = data_model/np.array([std_WX, std_y])
        data_obs = np.vstack([WX/std_WX, np.ravel(y)/std_y]).T
        t = tpoints[
            ((data_model[None, :, :] - data_obs[:, None, :]) ** 2).sum(axis=2).argmin(axis=1)
        ]
        return t

    def get_raw_velocity(self):
        self.velo_hat = use_model(self.WX_model, self.x, method=self.WX_method, bias=self.delta) - np.ravel(self.gamma * self.y)
        tmp1 = self.omega * self.t + self.theta
        self.y_t = self.alpha * np.sin(tmp1) + self.beta
        self.velo_t = self.alpha * self.omega * np.cos(tmp1)
        self.velo_normed = (self.velo_hat-self.velo_hat.min()) / (self.velo_hat.max()-self.velo_hat.min()+1e-6)
        return

default_pars_names = ["alpha", "beta", "omega", "theta", "gamma", "delta", "likelihood", "varx", "scaling_y"]
default_pars_names_vec = ["scaling", "weights", "weights_init", "weights_final"]
default_pars_names += default_pars_names_vec

def read_pars(adata, pars_names=None, key="fit"):
    pars = []
    for name in default_pars_names if pars_names is None else pars_names:
        pkey = f"{key}_{name}"
        if name in default_pars_names_vec:
            par = np.zeros([adata.n_vars, adata.var['n_TFs'].max()])
        else:
            par = np.zeros(adata.n_vars) * np.nan
        if pkey in adata.var.keys():
            par = adata.var[pkey].values
        if pkey in adata.varm.keys():
            par = adata.varm[pkey]
        pars.append(par)

    return pars

default_results_names = ["velo_hat", "velo_t", "y_t", "velo_normed", "filtered", "WX", "y"]
def write_result(adata, results):
    for i, name in enumerate(default_results_names):
        adata.layers[name] = results[i]

def write_pars(adata, pars, pars_names=None, add_key="fit"):
    for i, name in enumerate(default_pars_names if pars_names is None else pars_names):
        if name in default_pars_names_vec:
            adata.varm[f"{add_key}_{name}"] = pars[i]
        else:
            adata.var[f"{add_key}_{name}"] = pars[i]

def recover_dynamics(
    data,
    var_names="all",#"velocity_genes",
    n_top_genes=None,
    max_iter=10,
    WX_method='lsq_linear',
    WX_thres=20,
    max_n_TF=999,
    t_max=None,
    fit_time=True,
    fit_scaling=True,
    fit_steady_states=True,
    fit_connected_states=None,
    fit_basal_transcription=None,
    use_raw=False,
    load_pars=None,
    return_model=None,
    plot_results=False,
    steady_state_prior=None,
    add_key="fit",
    copy=False,
    n_jobs=None,
    backend="loky",
    reg=None,
    init_weight_method=None,
    fix_theta=0,
    gamma_thres=0,
    n_time_points=1000,
    **kwargs,
):
    """Recovers the full splicing kinetics of specified genes.

    The model infers transcription rates, splicing rates, degradation rates,
    as well as cell-specific latent time and transcriptional states,
    estimated iteratively by expectation-maximization.

    .. image:: https://user-images.githubusercontent.com/31883718/69636459-ef862800-1056-11ea-8803-0a787ede5ce9.png

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix.
    var_names: `str`,  list of `str` (default: `'velocity_genes'`)
        Names of variables/genes to use for the fitting. If `var_names='velocity_genes'`
        but there is no column `'velocity_genes'` in `adata.var`, velocity genes are
        estimated using the steady state model.
    n_top_genes: `int` or `None` (default: `None`)
        Number of top velocity genes to use for the dynamical model.
    max_iter:`int` (default: `10`)
        Maximal iterations in the EM-Algorithm.
    t_max: `float`, `False` or `None` (default: `None`)
        Total range for time assignments.
    fit_scaling: `bool` or `float` or `None` (default: `True`)
        Whether to fit scaling between unspliced and spliced.
    fit_time: `bool` or `float` or `None` (default: `True`)
        Whether to fit time or keep initially given time fixed.
    fit_steady_states: `bool` or `None` (default: `True`)
        Whether to explicitly model and fit steady states (next to induction/repression)
    fit_connected_states: `bool` or `None` (default: `None`)
        Restricts fitting to neighbors given by connectivities.
    fit_basal_transcription: `bool` or `None` (default: `None`)
        Enables model to incorporate basal transcriptions.
    use_raw: `bool` or `None` (default: `None`)
        if True, use .layers['sliced'], else use moments from .layers['Ms']
    load_pars: `bool` or `None` (default: `None`)
        Load parameters from past fits.
    return_model: `bool` or `None` (default: `None`)
        Whether to return the model as :DynamicsRecovery: object.
    plot_results: `bool` or `None` (default: `False`)
        Plot results after parameter inference.
    steady_state_prior: list of `bool` or `None` (default: `None`)
        Mask for indices used for steady state regression.
    add_key: `str` (default: `'fit'`)
        Key to add to parameter names, e.g. 'fit_t' for fitted time.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to `adata`.
    n_jobs: `int` or `None` (default: `None`)
        Number of parallel jobs.
    backend: `str` (default: "loky")
        Backend used for multiprocessing. See :class:`joblib.Parallel` for valid
        options.

    Returns
    -------
    fit_alpha: `.var`
        inferred transcription rates
    fit_beta: `.var`
        inferred splicing rates
    fit_gamma: `.var`
        inferred degradation rates
    fit_weight: `.var`
        inferred weight of a TF
    fit_t_: `.var`
        inferred switching time points
    fit_scaling: `.var`
        internal variance scaling factor for un/spliced counts
    fit_likelihood: `.var`
        likelihood of model fit
    fit_alignment_scaling: `.var`
        scaling factor to align gene-wise latent times to a universal latent time
    """  # noqa E501

    adata = data.copy() if copy else data

    n_jobs = get_n_jobs(n_jobs=n_jobs)
    logg.info(f"recovering dynamics (using {n_jobs}/{os.cpu_count()} cores)", r=True)

    if len(set(adata.var_names)) != len(adata.var_names):
        logg.warn("Duplicate var_names found. Making them unique.")
        adata.var_names_make_unique()

    if "M_total" not in adata.layers.keys():
        use_raw = True
    if fit_connected_states is None:
        fit_connected_states = not use_raw

    adata.uns["recover_dynamics"] = {
        "fit_connected_states": fit_connected_states,
        "fit_basal_transcription": fit_basal_transcription,
        "use_raw": use_raw,
    }

    if isinstance(var_names, str) and var_names not in adata.var_names:
        if var_names in adata.var.keys():
            if adata.var[var_names][0] in ['True', 'False']:
                adata.var[var_names] = adata.var[var_names].map({'True': True, 'False': False})
            var_names = adata.var_names[adata.var[var_names].values]
        elif use_raw or var_names == "all":
            var_names = adata.var_names
        elif "_genes" in var_names:
            from .velocity import Velocity

            velo = Velocity(adata, use_raw=use_raw)
            #velo.compute_deterministic(perc=[5, 95])
            #var_names = adata.var_names[velo._velocity_genes]
            #adata.var["fit_r2"] = velo._r2
        else:
            # raise ValueError("Variable name not found in var keys.")
            return False
    if not isinstance(var_names, str):
        var_names = list(np.ravel(var_names))

    var_names = make_unique_list(var_names, allow_array=True)
    var_names = np.array([name for name in var_names if name in adata.var_names])
    if len(var_names) == 0:
        # raise ValueError("Variable name not found in var keys.")
        return False
    if n_top_genes is not None and len(var_names) > n_top_genes:
        X = adata[:, var_names].layers[("total" if use_raw else "M_total")]
        var_names = var_names[np.argsort(np.sum(X, 0))[::-1][:n_top_genes]]
    if return_model is None:
        return_model = len(var_names) < 5

    pars = read_pars(adata)
    alpha, beta, omega, theta, gamma, delta, likelihood, varx, scaling_y = pars[:-4]
    scaling, weights, weights_init, weights_final = pars[-4], pars[-3], pars[-2], pars[-1]
    # likelihood[np.isnan(likelihood)] = 0
    idx, L = [], []
    velo_hat, velo_t, velo_normed = np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan
    y_t, y, filtered = np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan, np.zeros(adata.shape) * np.nan
    WX = np.zeros(adata.shape) * np.nan

    T = np.zeros(adata.shape) * np.nan
    if "fit_t" in adata.layers.keys():
        T = adata.layers["fit_t"]

    conn = get_connectivities(adata) if fit_connected_states else None

    res = parallelize(
        _fit_recovery,
        var_names,
        n_jobs,
        unit="gene",
        as_array=False,
        backend=backend,
        show_progress_bar=len(var_names) > 9,
    )(
        adata=adata,
        use_raw=use_raw,
        load_pars=load_pars,
        max_iter=max_iter,
        fit_time=fit_time,
        fit_steady_states=fit_steady_states,
        fit_scaling=fit_scaling,
        fit_basal_transcription=fit_basal_transcription,
        conn=conn,
        WX_method=WX_method,
        WX_thres=WX_thres,
        max_n_TF=max_n_TF,
        reg=reg,
        init_weight_method=init_weight_method,
        fix_theta=fix_theta,
        gamma_thres=gamma_thres,
        n_time_points=n_time_points,
        **kwargs,
    )
    idx, dms = map(_flatten, zip(*res)) # gene_id_list, model_list

    for ix, dm in zip(idx, dms):
        T[:, ix] = dm.t
        if WX_method in ['lsq_linear']:
            weight = np.ravel(dm.WX_model)
            weights[ix, :len(weight)] = weight
            weights_init[ix, :len(weight)] = np.ravel(dm.weights_init)
            weights_final[ix, :len(weight)] = np.ravel(dm.weights_final)
        alpha[ix], beta[ix], omega[ix], theta[ix], gamma[ix], delta[ix] = dm.alpha, dm.beta, dm.omega, dm.theta, dm.gamma, dm.delta #dm.pars[:5, -1]
        scaling[ix][:dm.n_TFs], scaling_y[ix] = dm.scaling, dm.scaling_y # dm.pars[10, -1],  dm.pars[7, -1]
        likelihood[ix], varx[ix] = dm.likelihood, dm.varx # dm.pars[5, -1],  dm.pars[6, -1]
        velo_hat[:,ix], velo_t[:,ix], velo_normed[:,ix] = dm.velo_hat.reshape(-1), dm.velo_t.reshape(-1), dm.velo_normed.reshape(-1)
        y_t[:,ix], y[:,ix], filtered[:,ix] =  dm.y_t.reshape(-1), dm.y.reshape(-1), dm.filtered.reshape(-1)
        WX[:,ix] = dm.WX.reshape(-1)
        L.append(dm.loss)

    _pars = [
        alpha, beta, omega, theta, gamma, delta, likelihood, varx, scaling_y, \
        scaling, weights, weights_init, weights_final
    ]
    write_pars(adata, _pars)
    write_result(adata, [velo_hat, velo_t, y_t, velo_normed, filtered, WX, y])

    adata.layers["fit_t_raw"] = T.copy()
    #if "fit_t" in adata.layers.keys():
    #    adata.layers["fit_t"][:, idx] = (
    #        T[:, idx] if conn is None else conn.dot(T[:, idx])
    #    )
    #else:
    adata.layers["fit_t"] = T if conn is None else conn.dot(T)

    if L:  # is False if only one invalid / irrecoverable gene was given in var_names
        cur_len = adata.varm["loss"].shape[1] if "loss" in adata.varm.keys() else 2
        max_len = max(np.max([len(loss) for loss in L]), cur_len) if L else cur_len
        loss = np.ones((adata.n_vars, max_len)) * np.nan

        if "loss" in adata.varm.keys():
            loss[:, :cur_len] = adata.varm["loss"]

        loss[idx] = np.vstack(
            [
                np.concatenate([loss, np.ones(max_len - len(loss)) * np.nan])
                for loss in L
            ]
        )
        adata.varm["loss"] = loss

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint(
        "added \n"
        f"    '{add_key}_pars', "
        f"fitted parameters for splicing dynamics (adata.var)"
    )

    '''
    P = []
    if plot_results:  # Plot Parameter Stats
        n_rows, n_cols = len(var_names[:4]), 6
        figsize = [2 * n_cols, 1.5 * n_rows]  # rcParams['figure.figsize']
        fontsize = rcParams["font.size"]
        fig, axes = pl.subplots(nrows=n_rows, ncols=6, figsize=figsize)
        pl.subplots_adjust(wspace=0.7, hspace=0.5)
        for i, gene in enumerate(var_names[:4]):
            #if t_max is not False:
            #    mi = dm.m[i]
            #    P[i] *= np.array([1 / mi, 1 / mi, 1 / mi, mi, 1])[:, None]
            ax = axes[i] if n_rows > 1 else axes
            for j, pij in enumerate(P[i]):
                ax[j].plot(pij)
            ax[len(P[i])].plot(L[i])
            if i == 0:
                pars_names = ["alpha", "beta", "gamma", "weight", "t_", "scaling", "loss"]
                for j, name in enumerate(pars_names):
                    ax[j].set_title(name, fontsize=fontsize)
    '''
    if return_model:
        logg.info("\noutputs model fit of gene:", dm.gene)

    #return dm if return_model else adata if copy else None
    return True

def latent_time(
    data,
    vkey="velocity",
    t_key='normed_t',
    min_likelihood=0.1,
    min_confidence=0.75,
    min_corr_diffusion=None,
    weight_diffusion=None,
    root_key=None,
    end_key=None,
    t_max=None,
    copy=False,
):
    """Computes a gene-shared latent time.

    Gene-specific latent timepoints obtained from the dynamical model are coupled to a
    universal gene-shared latent time, which represents the cell’s internal clock and
    is based only on its transcriptional dynamics.

    .. image:: https://user-images.githubusercontent.com/31883718/69636500-03318e80-1057-11ea-9e14-ae9f907711cc.png

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix
    vkey: `str` (default: `'velocity'`)
        Name of velocity estimates to be used.
    min_likelihood: `float` between `0` and `1` or `None` (default: `.1`)
        Minimal likelihood fitness for genes to be included to the weighting.
    min_confidence: `float` between `0` and `1` (default: `.75`)
        Parameter for local coherence selection.
    min_corr_diffusion: `float` between `0` and `1` or `None` (default: `None`)
        Only select genes that correlate with velocity pseudotime obtained
        from diffusion random walk on velocity graph.
    weight_diffusion: `float` or `None` (default: `None`)
        Weight applied to couple latent time with diffusion-based velocity pseudotime.
    root_key: `str` or `None` (default: `'root_cells'`)
        Key (.uns, .obs) of root cell to be used.
        If not set, it obtains root cells from velocity-inferred transition matrix.
    end_key: `str` or `None` (default: `None`)
        Key (.obs) of end points to be used.
    t_max: `float` or `None` (default: `None`)
        Overall duration of differentiation process.
        If not set, a overall transcriptional timescale of 20 hours is used as prior.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to `adata`.

    Returns
    -------
    latent_time: `.obs`
        latent time from learned dynamics for each cell
    """  # noqa E501

    adata = data.copy() if copy else data

    from .dynamical_model_utils import compute_shared_time, root_time
    from .terminal_states import terminal_states
    from .utils import scale, vcorrcoef
    from .velocity_graph import velocity_graph
    from .velocity_pseudotime import velocity_pseudotime

    if t_key not in adata.layers.keys():
        raise ValueError("you need to run `tl.recover_dynamics` first.")

    if f"{vkey}_graph" not in adata.uns.keys():
        velocity_graph(adata, approx=True)

    if root_key is None:
        terminal_keys = ["root_cells", "starting_cells", "root_states_probs"]
        keys = [key for key in terminal_keys if key in adata.obs.keys()]
        if len(keys) > 0:
            root_key = keys[0]
    if root_key not in adata.uns.keys() and root_key not in adata.obs.keys():
        root_key = "root_cells"
    if root_key not in adata.obs.keys():
        terminal_states(adata, vkey=vkey)

    t = np.array(adata.layers[t_key])
    idx_valid = ~np.isnan(t.sum(0))
    if min_likelihood is not None:
        likelihood = adata.var["fit_likelihood"].values
        idx_valid &= np.array(likelihood >= min_likelihood, dtype=bool)
    t = t[:, idx_valid]
    t_sum = np.sum(t, 1)
    conn = get_connectivities(adata)

    if root_key not in adata.uns.keys():
        roots = np.argsort(t_sum)
        idx_roots = np.array(adata.obs[root_key][roots])
        idx_roots[pd.isnull(idx_roots)] = 0
        if np.any([isinstance(ix, str) for ix in idx_roots]):
            idx_roots = np.array([isinstance(ix, str) for ix in idx_roots], dtype=int)
        idx_roots = idx_roots.astype(float) > 1 - 1e-3
        if np.sum(idx_roots) > 0:
            roots = roots[idx_roots]
        else:
            logg.warn(
                "No root cells detected. Consider specifying "
                "root cells to improve latent time prediction."
            )
    else:
        roots = [adata.uns[root_key]]
        root_key = f"root cell {adata.uns[root_key]}"

    if end_key in adata.obs.keys():
        fates = np.argsort(t_sum)[::-1]
        idx_fates = np.array(adata.obs[end_key][fates])
        idx_fates[pd.isnull(idx_fates)] = 0
        if np.any([isinstance(ix, str) for ix in idx_fates]):
            idx_fates = np.array([isinstance(ix, str) for ix in idx_fates], dtype=int)
        idx_fates = idx_fates.astype(float) > 1 - 1e-3
        if np.sum(idx_fates) > 0:
            fates = fates[idx_fates]
    else:
        fates = [None]

    logg.info(
        f"computing latent time using {root_key}"
        f"{', ' + end_key if end_key in adata.obs.keys() else ''} as prior",
        r=True,
    )

    VPT = velocity_pseudotime(
        adata, vkey, root_key=roots[0], end_key=fates[0], return_model=True
    )
    vpt = VPT.pseudotime

    if min_corr_diffusion is not None:
        corr = vcorrcoef(t.T, vpt)
        t = t[:, np.array(corr >= min_corr_diffusion, dtype=bool)]

    if root_key in adata.uns.keys():
        root = adata.uns[root_key]
        t, t_ = root_time(t, root=root)
        latent_time = compute_shared_time(t)
    else:
        roots = roots[:4]
        latent_time = np.ones(shape=(len(roots), adata.n_obs))
        for i, root in enumerate(roots):
            t, t_ = root_time(t, root=root)
            latent_time[i] = compute_shared_time(t)
        latent_time = scale(np.mean(latent_time, axis=0))

    if fates[0] is not None:
        fates = fates[:4]
        latent_time_ = np.ones(shape=(len(fates), adata.n_obs))
        for i, fate in enumerate(fates):
            t, t_ = root_time(t, root=fate)
            latent_time_[i] = 1 - compute_shared_time(t)
        latent_time = scale(latent_time + 0.2 * scale(np.mean(latent_time_, axis=0)))

    tl = latent_time
    tc = conn.dot(latent_time)

    z = tl.dot(tc) / tc.dot(tc)
    tl_conf = (1 - np.abs(tl / np.max(tl) - tc * z / np.max(tl))) ** 2
    idx_low_confidence = tl_conf < min_confidence

    if weight_diffusion is not None:
        w = weight_diffusion
        latent_time = (1 - w) * latent_time + w * vpt
        latent_time[idx_low_confidence] = vpt[idx_low_confidence]
    else:
        conn_new = conn.copy()
        conn_new[:, idx_low_confidence] = 0
        conn_new.eliminate_zeros()
        latent_time = conn_new.dot(latent_time)

    latent_time = scale(latent_time)
    if t_max is not None:
        latent_time *= t_max

    adata.obs["latent_time"] = latent_time

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint("added \n" "    'latent_time', shared time (adata.obs)")
    return adata if copy else None

recover_latent_time = latent_time

def rank_dynamical_genes(data, n_genes=100, groupby=None, copy=False):
    """Rank genes by likelihoods per cluster/regime.

    This ranks genes by their likelihood obtained from the
    dynamical model grouped by clusters specified in groupby.

    Arguments
    ----------
    data : :class:`~anndata.AnnData`
        Annotated data matrix.
    n_genes : `int`, optional (default: 100)
        The number of genes that appear in the returned tables.
    groupby: `str`, `list` or `np.ndarray` (default: `None`)
        Key of observations grouping to consider.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to data.

    Returns
    -------
    rank_dynamical_genes : `.uns`
        Structured array to be indexed by group id storing the gene
        names. Ordered according to scores.
    """

    from .dynamical_model_utils import get_divergence

    adata = data.copy() if copy else data

    logg.info("ranking genes by cluster-specific likelihoods", r=True)

    groupby = (
        groupby
        if isinstance(groupby, str) and groupby in adata.obs.keys()
        else "clusters"
        if "clusters" in adata.obs.keys()
        else "louvain"
        if "louvain" in adata.obs.keys()
        else "velocity_clusters"
        if "velocity_clusters" in adata.obs.keys()
        else None
    )

    vdata = adata[:, ~np.isnan(adata.var["fit_alpha"])]
    groups = vdata.obs[groupby].cat.categories

    ll = get_divergence(
        vdata,
        mode="gene_likelihood",
        use_connectivities=True,
        clusters=adata.obs[groupby],
    )
    idx_sorted = np.argsort(np.nan_to_num(ll), 1)[:, ::-1][:, :n_genes]
    rankings_gene_names = vdata.var_names[idx_sorted]
    rankings_gene_scores = np.sort(np.nan_to_num(ll), 1)[:, ::-1][:, :n_genes]

    key = "rank_dynamical_genes"
    if key not in adata.uns.keys():
        adata.uns[key] = {}

    adata.uns[key] = {
        "names": np.rec.fromarrays(
            [n for n in rankings_gene_names],
            dtype=[(f"{rn}", "U50") for rn in groups],
        ),
        "scores": np.rec.fromarrays(
            [n.round(2) for n in rankings_gene_scores],
            dtype=[(f"{rn}", "float32") for rn in groups],
        ),
    }

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint("added \n" f"    '{key}', sorted scores by group ids (adata.uns)")

    return adata if copy else None

def _fit_recovery(
    var_names,
    adata,
    use_raw,
    load_pars,
    max_iter,
    fit_time,
    fit_steady_states,
    conn,
    fit_scaling,
    fit_basal_transcription,
    max_n_TF,
    queue,
    WX_method,
    reg,
    WX_thres,
    init_weight_method,
    fix_theta,
    gamma_thres,
    n_time_points,
    **kwargs,
):

    idx, dms = [], []
    for gene in var_names:#################
        dm = DynamicsRecovery(
            adata,
            gene,
            use_raw=use_raw,
            load_pars=load_pars,
            max_iter=max_iter,
            fit_time=fit_time,
            fit_steady_states=fit_steady_states,
            fit_connected_states=conn,
            fit_scaling=fit_scaling,
            fit_basal_transcription=fit_basal_transcription,
            WX_method=WX_method,
            WX_thres=WX_thres,
            max_n_TF=max_n_TF,
            reg=reg,
            init_weight_method=init_weight_method,
            fix_theta=fix_theta,
            gamma_thres=gamma_thres,
            n_time_points=n_time_points,
            **kwargs,
        )
        tmp = str(list(var_names).index(gene)) + '/' + str(len(var_names))
        if dm.recoverable:
            #try:
            print('Processing', tmp, dm.gene)
            dm.fit()

            ix = np.where(adata.var_names == gene)[0][0]
            idx.append(ix)
            dms.append(dm)
            print(tmp, dm.gene, 'FINISHED with n_TFs:', dm.n_TFs)
            #except:
            #    dm.recoverable = False
            #    logg.warn(dm.gene, "not recoverable during optimization.")

        else:
            logg.warn(tmp, dm.gene, "not recoverable due to insufficient samples.")

        if queue is not None:
            queue.put(1)

    if queue is not None:
        queue.put(None)

    return idx, dms # gene_id_list, model_list

def _flatten(iterable):
    return [i for it in iterable for i in it]


import os.path
import random
from pprint import pprint
import networkx as nx
import numpy as np

import torch
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score
from torch import optim as optim
from tqdm import tqdm
from torch.nn.utils.rnn import pad_sequence

from load import load_video_id
from data import LMDBFile
from load import load_query_input
from load import load_concepts
from model import DeepQNetIVCML
from env import IVCMLEnv
from util import cuda
from vocab import Vocab
from loss import TripletLoss

from collections import OrderedDict
from util import mean
from util import flatten
import json
from util import MAX_MOMENT_DISTANCE

class Trainer(object):
    def __init__(self, opts):
        super(Trainer, self).__init__()
        self.opts = opts
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        self.feature_type = opts.feature_type

        # Load concepts
        self.data_root = opts.data_root
        self.feature_root = opts.feature_root
        self.dataset = opts.dataset
        concept_list_path = os.path.join(self.feature_root, f"{self.dataset}_dual_task_emb_concept_p3_2048_concept_bank.txt")
        concepts, concept_size = load_concepts(concept_list_path)
        # Load video id
        video_id_path = os.path.join(self.feature_root, f"{self.dataset}_dual_task_emb_concept_p3_2048_vid_id.txt")
        shot_ids, shot_id_size = load_video_id(video_id_path)
        self.vocab = Vocab(concepts, opts.stem)

        self.env = IVCMLEnv(self.opts, concepts, shot_ids)
        self.minimum_distance = 1
        self.maximum_distance = 4
        self.infer_steps = self.opts.infer_steps
        self.dqn = DeepQNetIVCML(env=self.env,
                                 d_fea=768,
                                 d_concept_en=self.vocab.concept_size_en,
                                 d_concept_de=self.vocab.concept_size_de,
                                 infer_steps=self.infer_steps,
                                 opts=opts)
        self.dqn.to(self.device)
        self.optimizer = optim.Adam(self.dqn.parameters(), lr=opts.init_lr)
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=1, gamma=0.8)  # step based on conditions

        self.checkpoint = None
        self.is_inference = False
        self.cur_data_size = 0

        # display training info
        pprint(vars(opts))

        # General configuration
        self.training_prefix = opts.training_prefix
        self.judge_metric = "acc"
        self.best_model_path = "%s/%s_best_%s" % (opts.ckpt_root, self.training_prefix, self.judge_metric)
        print("Best checkpoint will be stored at: %s." % self.best_model_path)

        # Training configuration
        self.num_epoch = 200
        self.batch_size = opts.batch_size  # 4 for nk3
        self.evaluate_period = 1
        self.ent_coef = opts.ent_coef
        self.vf_coef = opts.vf_coef
        self.metric_cats = ["acc", "recall", "precision", "f1"]
        self.early_stop_epoch_num = opts.early_stop_epoch_num
        self.transition = opts.transition
        self.normalize_advantage = True
        self.loss_weight = [float(w) for w in opts.loss_weight.split(",")]
        self.loss_func = opts.loss_func.split(",")
        for func in self.loss_func:
            assert func in ["a2c", "ce", "triplet"], "Error: Unsupported loss function: %s" % func
            print("%s weight: %f" % (func, self.loss_weight[self.loss_func.index(func)]))
        self.select_by = opts.select_by
        self.max_grad_norm = 2

        # Inference configuration
        self.step_used = {}
        self.traj_collection = {}
        self.traj_degree = []
        self.c_recorded_sample = 0

        # recorder
        self.ce_loss_epoch = []
        self.triplet_loss_epoch = []
        self.policy_loss_epoch = []
        self.value_loss_epoch = []
        self.entropy_loss_epoch = []
        self.c_covered_step_nk, self.c_out_step_nk = 0, 0

        # global variables
        self.no_improve_epoch_num = 0
        self.is_early_stop = False
        self.no_decay_epoch_num = 0
        self.is_decay = False
        self.best_metrics = {"acc": 0, "recall": 0, "precision": 0, "f1": 0}
        self.i_epoch = 0

        # Loss functions
        self.ce_loss_func = torch.nn.CrossEntropyLoss()
        self.mse_loss_func = torch.nn.MSELoss()
        self.triplet_loss_func = TripletLoss(margin=self.opts.margin,
                                             measure="cosine",
                                             max_violation=False,
                                             cost_style="all",
                                             direction="all",
                                             seq_match=False,
                                             pos_loss_term=False)

    def run(self):
        input_reader = self.load_input_reader()
        if self.opts.mode in ["train", "infer"]:
            sample_ids_trn, sample_ids_val, sample_ids_tst = self.get_sample_partition(input_reader)
            print("Train: %d; Validation: %d; Test: %d" % (len(sample_ids_trn), len(sample_ids_val), len(sample_ids_tst)))
        elif self.opts.mode in ["infer_hero", "infer_conquer"]:
            sample_ids_trn, sample_ids_val = None, None
            sample_ids_tst = self.get_real_inference_sample(input_reader)
        else:
            assert False

        if self.opts.mode == "train":
            self.train_from_scratch(input_reader, sample_ids_trn, sample_ids_val, sample_ids_tst)
        elif self.opts.mode in ["infer", "infer_hero", "infer_conquer"]:
            self.infer(input_reader, sample_ids_tst)
        else:
            print("Error: Invalid mode '%s'" % self.opts.mode)
            raise NotImplementedError

    def train_from_scratch(self, input_reader, sample_ids_trn, sample_ids_val, sample_ids_tst):
        if self.opts.resume:
            self.load_model(self.best_model_path)
            training_range = range(self.checkpoint["epoch"] + 1, self.num_epoch)
            print(f"Resume training from Epoch {self.checkpoint['epoch'] + 1}.")
        else:
            training_range = range(self.num_epoch)

        if self.opts.vr_initialize:
            assert self.opts.task == "ivcml", "Error: cannot initialize model by pretrained vr in %s task" % self.opts.task
            self.load_vr_pretrain()

        for i_epoch in training_range:
            self.i_epoch = i_epoch
            self.train_epoch(input_reader, sample_ids_trn)
            self.env.summary()
            self.env.clear_cache()
            if i_epoch % self.evaluate_period == 0:
                metric_epoch = self.evaluate(input_reader, sample_ids_val, mode="val")
                self.env.summary()
                self.env.clear_cache()
                self.judge(metric_epoch)
            if self.is_early_stop:
                print("Early stop happens.")
                break
            if self.is_decay:
                self.scheduler.step()
            print("Learning rate @ E-%d:" % self.i_epoch, self.scheduler.get_last_lr())

            self.clear_cache()

        # Test model
        self.load_model(self.best_model_path)
        metric_test = self.evaluate(input_reader, sample_ids_tst, mode="test")
        self.env.summary()
        self.env.clear_cache()

    def infer(self, input_reader, sample_ids_tst):
        self.is_inference = True
        random_init = self.opts.random_policy
        if random_init:
            print("Randomly initialize without loading the trained model.")
        else:
            self.load_model(self.best_model_path)

        print("Evaluate on test:")
        self.evaluate(input_reader, sample_ids_tst, mode="test")
        self.clear_cache()

    def batch_generator(self, input_reader, sample_ids, is_train=True):
        data_size = len(sample_ids)
        batch_num = data_size // self.batch_size + int((data_size % self.batch_size) != 0)

        c_sample = 0
        with tqdm(range(batch_num), total=batch_num) as t:
            for i in t:
                st = i * self.batch_size
                ed = st + self.batch_size
                batch_sample_ids = sample_ids[st: ed]
                c_sample += len(batch_sample_ids)
                query_id_batch = []
                targets_batch, start_batch = [], []
                query_text_batch = []
                input_ids_bert_batch, input_mask_bert_batch, segment_ids_bert_batch = [], [], []
                proposal_batch, distance_proposal_batch = [], []
                bow_vec_batch = []
                tok_ids_batch = []
                query_rank_batch = []
                init_distance_batch = []
                for sample_id in batch_sample_ids:
                    query_id, start_id, node_ids, query_text, \
                        input_ids_bert, input_mask_bert, segment_ids_bert, \
                        bow_vec, tok_ids, query_rank, init_distance, \
                        proposal, distance_proposal = load_query_input(sample_id, input_reader, self.vocab)
                    query_text_batch.append(query_text)
                    input_ids_bert_batch.append(input_ids_bert)
                    input_mask_bert_batch.append(input_mask_bert)
                    segment_ids_bert_batch.append(segment_ids_bert)
                    query_id_batch.append(query_id)
                    start_batch.append(start_id)
                    targets_batch.append(node_ids)
                    bow_vec_batch.append(bow_vec)
                    tok_ids_batch.append(tok_ids)
                    query_rank_batch.append(query_rank)
                    init_distance_batch.append(init_distance)
                    proposal_batch.append(proposal)
                    distance_proposal_batch.append(distance_proposal)

                input_ids_bert_batch = pad_sequence(input_ids_bert_batch, batch_first=True)
                input_mask_bert_batch = pad_sequence(input_mask_bert_batch, batch_first=True)
                segment_ids_bert_batch = pad_sequence(segment_ids_bert_batch, batch_first=True)
                tok_ids_batch = pad_sequence(tok_ids_batch, batch_first=True)
                bow_vec_batch = torch.tensor(np.concatenate(bow_vec_batch, axis=0), dtype=torch.float32)

                targets_batch = pad_sequence(targets_batch, batch_first=True, padding_value=-2)  # default final pred is -1
                start_batch = torch.tensor(start_batch, dtype=torch.long)
                query_id_batch = torch.tensor(query_id_batch, dtype=torch.long)
                query_rank_batch = torch.tensor(query_rank_batch, dtype=torch.long)
                init_distance_batch = torch.tensor(init_distance_batch, dtype=torch.long)
                proposal_batch = torch.tensor(proposal_batch, dtype=torch.long)
                distance_proposal_batch = torch.tensor(distance_proposal_batch, dtype=torch.long)

                if is_train:
                    ce_loss, tri_loss, value_loss, pi_loss, pi_entropy = self.batch_loss_value()

                    info = f"E:{self.i_epoch}|"
                    info_terms = [(f"π:{pi_loss :.3f}|", pi_loss),
                                  (f"v:{value_loss :.3f}|", value_loss),
                                  (f"ce:{ce_loss :.3f}|", ce_loss),
                                  (f"tri:{tri_loss :.3f}|", tri_loss),
                                  (f"H:{pi_entropy :.3f}|", pi_entropy)]

                    for info_term, v in info_terms:
                        if v:
                            info += info_term

                    info += f"BA:{self.best_metrics['acc']:.3f}|" \
                            f"NIE:{self.no_improve_epoch_num}/{self.early_stop_epoch_num}|" \
                            f"mf:{self.env.c_missed_node}/{len(self.env.missed_feature_id)}"

                    t.set_description(desc=info)

                yield i, (input_ids_bert_batch, input_mask_bert_batch, segment_ids_bert_batch, tok_ids_batch, bow_vec_batch, start_batch, targets_batch, query_id_batch, query_rank_batch,
                          init_distance_batch, proposal_batch, distance_proposal_batch, query_text_batch)
        self.cur_data_size = 0
        print("Generated %d samples in total" % c_sample)

    def iterate(self, batch):
        input_ids_bert_batch, input_mask_bert_batch, segment_ids_bert_batch, tok_ids_batch, bow_vec_batch, \
            start_batch, targets_batch, query_id_batch, query_rank_batch, init_distance_batch, \
            proposal, distance_proposal, query_text_batch = cuda(batch, self.device)

        outputs = self.dqn(input_ids_bert_batch, input_mask_bert_batch, segment_ids_bert_batch, tok_ids_batch, bow_vec_batch, start_batch, query_id_batch, targets_batch)
        return (*outputs, query_rank_batch, init_distance_batch, start_batch, targets_batch)

    def train_epoch(self, input_reader, sample_ids_trn):
        self.dqn.train()
        self.dqn.training = True
        self.env.train_mode()

        for i_batch, batch_data in self.batch_generator(input_reader, sample_ids_trn, is_train=True):
            policy_pred_batch, value_pred_batch, last_value_batch, \
                policy_gold_batch, gae_batch, returns_batch, \
                policy_mask_batch, step_mask_batch, states_batch, concept_steps, \
                target_batch, query_ids, query_rank, init_distance, start_batch_verify, targets_batch_verify = self.iterate(batch_data)

            # index valid steps
            valid_step_mask = step_mask_batch[:, :-1].nonzero(as_tuple=True)  # step_mask_batch includes last state.
            policy_pred_batch = policy_pred_batch[valid_step_mask]
            value_pred_batch = value_pred_batch[valid_step_mask]
            gae_batch = gae_batch[valid_step_mask]
            returns_batch = returns_batch[valid_step_mask]
            policy_mask_batch = policy_mask_batch[valid_step_mask]
            policy_gold_batch = policy_gold_batch[valid_step_mask]

            if self.normalize_advantage:
                gae_batch = (gae_batch - gae_batch.mean()) / (gae_batch.std() + 1e-8)

            mask = policy_gold_batch >= 0
            self.c_covered_step_nk += mask.sum().item()
            self.c_out_step_nk += (policy_gold_batch < 0).sum().item()
            coo_covered = mask.nonzero(as_tuple=True)
            policy_pred_batch, policy_mask_batch, policy_gold_batch = policy_pred_batch[coo_covered], policy_mask_batch[coo_covered], policy_gold_batch[coo_covered]
            value_pred_batch, gae_batch, returns_batch = value_pred_batch[coo_covered], gae_batch[coo_covered], returns_batch[coo_covered]

            # back propagation
            loss = self.calc_loss(policy_pred_batch, policy_mask_batch, policy_gold_batch, value_pred_batch, gae_batch, returns_batch)
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.dqn.parameters(), self.max_grad_norm)
            self.optimizer.step()

    @torch.no_grad()
    def evaluate(self, input_reader, query_ids_evaluate, mode="val"):
        self.dqn.eval()
        self.dqn.training = False
        self.env.eval_mode()

        policy_pred_eval, policy_gold_eval, value_pred_eval = [], [], []
        num_step_used = []
        for i_batch, batch_data in self.batch_generator(input_reader, query_ids_evaluate, is_train=True):
            policy_pred_batch, value_pred_batch, last_value_batch, \
                policy_gold_batch, gae_batch, returns_batch, \
                neighbor_mask_batch, step_mask_batch, states_batch, concept_steps, \
                target_batch, query_ids, query_rank, init_distance, start_batch_verify, targets_batch_verify = self.iterate(batch_data)

            # bad variable name: the initialized states at each step and,
            # meanwhile, they are also the predicted states in last time step.
            distance_batch = self.env.get_all_steps_distance_to_moment(states_batch.detach().cpu(), query_ids)

            # drop start to calculate number of steps used
            num_step_used_batch = step_mask_batch.sum(dim=-1)
            num_step_used.append(num_step_used_batch)

            self.record_and_analyze_inference(states_batch, concept_steps, target_batch, query_ids, query_rank, distance_batch, init_distance, num_step_used_batch)

            # drop the last to index the valid samples
            step_mask_init_batch = step_mask_batch[:, :-1]
            valid_step_mask = step_mask_init_batch.nonzero(as_tuple=True)
            policy_pred_batch = policy_pred_batch[valid_step_mask].detach().cpu()
            neighbor_mask_batch = neighbor_mask_batch[valid_step_mask].detach().cpu()
            policy_gold_batch = policy_gold_batch[valid_step_mask].detach().cpu()

            invalid_nei_coo = torch.logical_not(neighbor_mask_batch).nonzero(as_tuple=True)
            policy_pred_batch[invalid_nei_coo] = -1e5

            policy_pred_batch = policy_pred_batch.argmax(dim=-1)
            policy_gold_eval.append(policy_gold_batch)
            policy_pred_eval.append(policy_pred_batch)

        print("Overall Step-wise Accuracy:")
        print(self.opts.training_prefix, "\n****** %s ******" % self.opts.mode)
        metrics = self.calc_path_metric(num_step_used, mode)
        self.summary()
        return metrics

    def calc_metrics(self, gold_eval, pred_eval, mode):
        if pred_eval.sum():
            acc = accuracy_score(gold_eval, pred_eval)
            recall = recall_score(gold_eval, pred_eval)
            precision = precision_score(gold_eval, pred_eval)
            f1 = f1_score(gold_eval, pred_eval)
        else:
            acc, recall, precision, f1 = -1, -1, -1, -1

        metrics = {
            "acc": acc,
            "recall": recall,
            "precision": precision,
            "f1": f1
        }

        if mode == "val":
            res_info = f"Acc: {acc:.6f} | " \
                       f"Best Acc: {max(acc, self.best_metrics['acc']):.6f} | " \
                       f"NIE: {self.no_improve_epoch_num}/{self.early_stop_epoch_num}"
        else:
            res_info = f"Acc: {acc:.6f}"

        info = {
            "val": "Epoch: %d Validation %s " % (self.i_epoch, res_info),
            "test": "Test %s" % res_info
        }[mode]
        print(info)

        return metrics

    def calc_step_metric(self, gold_list, pred_list, mode):
        print("Step Accuracy:", end=" ")
        pred = torch.cat(gold_list, dim=0)
        gold = torch.cat(pred_list, dim=0)
        pred = (pred == gold).to(torch.long).detach().cpu()
        gold = torch.ones_like(pred, dtype=torch.long, device=torch.device("cpu"))
        metrics = self.calc_metrics(gold, pred, mode)
        return metrics

    def calc_path_metric(self, num_step_used, mode):
        num_step_used = torch.cat(num_step_used, dim=0)
        pred = (num_step_used <= self.infer_steps).to(torch.long).detach().cpu()
        gold = torch.ones_like(pred, dtype=torch.long, device=torch.device("cpu"))
        n_true = pred.sum().item()
        n_false = (1 - pred).sum().item()
        ttl = n_true + n_false
        print("Path Accuracy: %d/%d | False: %d/%d |" % (n_true, ttl, n_false, ttl), end=" ")
        metrics = self.calc_metrics(gold, pred, mode)
        return metrics

    def judge(self, metrics):
        best_metrics_upd = {}
        for key in self.metric_cats:
            best_metrics_upd[key] = max(metrics[key], self.best_metrics[key])
            if key != self.judge_metric:
                continue
            if metrics[key] > self.best_metrics[key] + 1e-5:
                self.no_improve_epoch_num = 0
                self.save_model(metrics)
            else:
                self.no_improve_epoch_num += 1

            if self.no_improve_epoch_num >= 5 and self.no_decay_epoch_num >= 5:
                self.is_decay = True
                self.no_decay_epoch_num = 0
            else:
                self.no_decay_epoch_num += 1
                self.is_decay = False

        early_stop = self.no_improve_epoch_num == self.early_stop_epoch_num
        self.best_metrics = best_metrics_upd
        self.is_early_stop = early_stop

    def save_model(self, metrics):
        ce_loss, tri_loss, value_loss, policy_loss, entropy_loss = self.epoch_loss_value()

        checkpoint = {
            'epoch': self.i_epoch,
            'model_state_dict': self.dqn.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'ce_loss': ce_loss,
            'tri_loss': tri_loss,
            'value_loss': value_loss,
            'policy_loss': policy_loss,
            'entropy_loss': entropy_loss,
        }
        checkpoint.update(metrics)
        torch.save(checkpoint, self.best_model_path)
        print("Model '%s' saved!" % self.best_model_path)

        del checkpoint["model_state_dict"]
        del checkpoint["optimizer_state_dict"]
        pprint(checkpoint)

    def load_model(self, model_path):
        checkpoint = torch.load(model_path, self.device)
        self.dqn.load_state_dict(checkpoint["model_state_dict"])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

        del checkpoint["model_state_dict"]
        del checkpoint["optimizer_state_dict"]
        print("Model '%s' loaded!" % model_path)
        self.checkpoint = checkpoint
        pprint(checkpoint)

    def load_vr_pretrain(self):
        if self.opts.dataset == "didemo":
            model_path = os.path.join(self.opts.ckpt_root, f"{self.opts.dataset}_vr_fbert_wms_gru_best_r_match")
        elif self.opts.dataset == "tvr":
            model_path = os.path.join(self.opts.ckpt_root, f"{self.opts.dataset}_vr_fbert_gru_best_r_match")
        else:
            print("Unsupported dataset %s" % self.opts.dataset)
            raise NotImplementedError

        vr_checkpoint = torch.load(model_path, self.device)
        vr_param_keys = vr_checkpoint["model_state_dict"].keys()
        param_keys = self.dqn.state_dict().keys()
        shared_keys = [k for k in param_keys if k in vr_param_keys]
        state_dict = OrderedDict([(k, vr_checkpoint["model_state_dict"][k] if k in vr_param_keys else self.dqn.state_dict()[k]) for k in param_keys])
        self.dqn.load_state_dict(state_dict)
        print("VR Model '%s' loaded!" % model_path)
        print("The following parameters (%d) are loaded:" % len(shared_keys))
        pprint(sorted(shared_keys, reverse=True))
        del vr_checkpoint["model_state_dict"]
        del vr_checkpoint["optimizer_state_dict"]
        pprint(vr_checkpoint)

    def clear_cache(self):
        self.c_covered_step_nk = 0
        self.c_out_step_nk = 0

        self.step_used = {}
        self.traj_collection = {}
        self.traj_degree = []
        self.c_recorded_sample = 0

        self.ce_loss_epoch = []
        self.policy_loss_epoch = []
        self.value_loss_epoch = []
        self.entropy_loss_epoch = []

    # loss
    def calc_loss(self, policy_pred_batch, policy_mask_batch, policy_gold_batch, value_pred_batch, gae_batch, returns_batch):
        loss = []
        if "a2c" in self.loss_func:
            a2c_weight = self.loss_weight[self.loss_func.index("a2c")]
            a2c_loss = self.calc_a2c_loss(policy_pred_batch, value_pred_batch, gae_batch, returns_batch, a2c_weight)
            loss.append(a2c_loss)

        if "ce" in self.loss_func:
            ce_weight = self.loss_weight[self.loss_func.index("ce")]
            ce_loss = ce_weight * self.ce_loss_func(policy_pred_batch, policy_gold_batch)
            self.ce_loss_epoch.append(ce_loss.item())
            loss.append(ce_loss)

        if "triplet" in self.loss_func:
            triplet_weight = self.loss_weight[self.loss_func.index("triplet")]
            tri_loss = triplet_weight * self.triplet_loss_func.ranking_loss_one_to_more(policy_pred_batch, policy_gold_batch, policy_mask_batch)
            self.triplet_loss_epoch.append(tri_loss.item())
            loss.append(tri_loss)

        loss = sum(loss)
        return loss

    def calc_a2c_loss(self, policy_pred_batch, value_pred_batch, advantage, returns, weight=1.0):
        # Gradient Loss
        pi = F.softmax(policy_pred_batch, dim=-1).max(dim=-1, keepdim=False).values
        neg_log_prob = - torch.log(pi)
        policy_entropy = - torch.mean(neg_log_prob * pi)
        policy_loss = (neg_log_prob * advantage).mean()
        # Value Loss
        value_loss = self.mse_loss_func(value_pred_batch, returns)
        a2c_loss = policy_loss + self.ent_coef * policy_entropy + self.vf_coef * value_loss

        self.policy_loss_epoch.append(weight*policy_loss.item())
        self.value_loss_epoch.append(weight*self.vf_coef*value_loss.item())
        self.entropy_loss_epoch.append(weight*self.ent_coef*policy_entropy.item())
        return weight*a2c_loss

    def load_input_reader(self):
        tag_graph_param = f"wc_{self.opts.weight_concept:.2f}_ws_{self.opts.weight_subtitle:.2f}_ths_{self.opts.edge_threshold_single:.2f}_thx_{self.opts.edge_threshold_cross:.2f}"
        if self.opts.use_balanced_graph:
            tag_graph_param += f"_top_{self.opts.k_ei}"
        if self.opts.mode in ["train", "infer"]:
            input_tag = f"clip_level_input_{tag_graph_param}"
        elif self.opts.mode in ["infer_hero"]:
            input_tag = f"clip_level_input_hero_{tag_graph_param}"
        elif self.opts.mode in ["infer_conquer"]:
            input_tag = f"clip_level_input_conquer_{tag_graph_param}"
        else:
            assert False

        input_reader = LMDBFile(self.opts.data_root, self.feature_type, self.opts.split, tag=input_tag, encode_method="json", readonly=True)
        return input_reader

    def get_sample_partition(self, input_reader):
        """ Sample partition for training, validation and testing """

        # parsing from sample id
        def __sample_id_to_distance(s_id):
            return int(s_id.split("_")[2])

        def __sample_id_to_query_id(s_id):
            return s_id.split("_")[0]

        def __sample_id_to_split(s_id):
            return s_id.split("_")[3]

        sample_ids = input_reader.keys()
        sample_id_split = {"train": [], "val": [], "test": []}
        split_c_legal_d = {"train": 0, "val": 0, "test": 0}
        split_c_legal_n = {"train": 0, "val": 0, "test": 0}
        cc_size_minimum = 600

        # prepare connected component (cc) size for filtering
        node_cc_size = [-1 for _ in self.env.graph.nodes()]
        ccs = nx.algorithms.components.connected_components(self.env.graph)
        for cc in ccs:
            cc_size = len(cc)
            for n in cc:
                node_cc_size[n] = cc_size

        # process training samples
        for sample_id in tqdm(sample_ids, total=len(sample_ids), desc="Splitting..."):
            split = __sample_id_to_split(sample_id)

            # 1. filter by initial distance
            is_legal_d = self.minimum_distance <= __sample_id_to_distance(sample_id) <= self.maximum_distance
            if not is_legal_d:
                continue
            split_c_legal_d[split] += 1

            # 2. filter by connected component size
            tgt_probe = input_reader[sample_id]["target"][0]
            cc_size = node_cc_size[tgt_probe]
            if cc_size < cc_size_minimum:
                continue
            split_c_legal_n[split] += 1

            # split valid samples
            sample_id_split[split].append(sample_id)

        for split in split_c_legal_d:
            print("%s: Get %d/%d with cc size > %d." % (split, split_c_legal_n[split], split_c_legal_d[split], cc_size_minimum))

        if sample_id_split["test"]:
            return sample_id_split["train"], sample_id_split["val"], sample_id_split["test"]

        # Split training set into train and validation in case of no testing set
        id_by_tgt = {}
        for sample_id in sample_id_split["train"]:
            tgt = __sample_id_to_query_id(sample_id)
            if tgt not in id_by_tgt:
                id_by_tgt[tgt] = []
            id_by_tgt[tgt].append(sample_id)

        target_ids = list(id_by_tgt.keys())
        random.seed(42)
        random.shuffle(target_ids)

        query_size_ttl = len(target_ids)
        boundary = int(query_size_ttl * 0.8)
        target_ids_trn = target_ids[:boundary]
        target_ids_val = target_ids[boundary:]

        sample_ids_trn = flatten([id_by_tgt[tgt] for tgt in target_ids_trn])
        sample_ids_val = flatten([id_by_tgt[tgt] for tgt in target_ids_val])
        sample_ids_test = sample_id_split["val"]
        return sample_ids_trn, sample_ids_val, sample_ids_test

    def get_real_inference_sample(self, input_reader):
        sample_ids = input_reader.keys()
        sample_ids_tst = []
        inference_split = "val" if self.opts.dataset == "tvr" else "test"
        for sample_id in tqdm(sample_ids, total=len(sample_ids), desc="Splitting..."):
            if input_reader[sample_id]["split"] != inference_split:
                continue
            if int(input_reader[sample_id]["distance"]) < MAX_MOMENT_DISTANCE:  # update distance
                distance = int(input_reader[sample_id]["distance"])
                self.maximum_distance = max(self.maximum_distance, distance)
                self.minimum_distance = min(self.minimum_distance, distance)
            sample_ids_tst.append(sample_id)
        return sample_ids_tst

    # display functions
    def epoch_loss_value(self):
        ce_loss = mean(self.ce_loss_epoch) if self.ce_loss_epoch else 0
        tri_loss = mean(self.triplet_loss_epoch) if self.triplet_loss_epoch else 0
        value_loss = mean(self.value_loss_epoch) if self.value_loss_epoch else 0
        policy_loss = mean(self.policy_loss_epoch) if self.policy_loss_epoch else 0
        entropy_loss = mean(self.entropy_loss_epoch) if self.entropy_loss_epoch else 0
        return ce_loss, tri_loss, value_loss, policy_loss, entropy_loss

    def batch_loss_value(self):
        ce_loss = self.ce_loss_epoch[-1] if self.ce_loss_epoch else 0
        tri_loss = self.triplet_loss_epoch[-1] if self.triplet_loss_epoch else 0
        value_loss = self.value_loss_epoch[-1] if self.value_loss_epoch else 0
        policy_loss = self.policy_loss_epoch[-1] if self.policy_loss_epoch else 0
        policy_entropy = self.entropy_loss_epoch[-1] if self.entropy_loss_epoch else 0
        return ce_loss, tri_loss, value_loss, policy_loss, policy_entropy

    def record_and_analyze_inference(self, states_with_last: torch.Tensor, concept_steps: torch.Tensor,
                                     target_ids: torch.Tensor, query_ids: torch.Tensor, query_rank: torch.Tensor,
                                     state_distance: torch.Tensor, init_distance: torch.Tensor, step_used: torch.Tensor):

        batch_size = states_with_last.shape[0]
        initial_distance = state_distance[:, 0].tolist()
        target_ids = target_ids.detach().cpu()

        for i_sample in range(batch_size):
            init_d = initial_distance[i_sample]
            tgt_ids = target_ids[i_sample].tolist()
            query_id = query_ids[i_sample].item()
            rank = query_rank[i_sample].item()
            if init_d not in self.step_used:
                self.step_used[init_d] = {}

            num_step = step_used[i_sample].item()
            if num_step not in self.step_used[init_d]:
                self.step_used[init_d][num_step] = 0

            self.step_used[init_d][num_step] += 1

            states_tra = states_with_last[i_sample][:num_step + 1].tolist()
            distances_tra = state_distance[i_sample][:num_step + 1].tolist()
            if not self.opts.no_feedback and not self.opts.random_policy:
                concepts = concept_steps[i_sample][:num_step + 1].tolist()
            else:
                concepts = []

            key = "d%ds%d" % (init_d, num_step)
            if key not in self.traj_collection:
                self.traj_collection[key] = {}
            query_key = "%d_%d" % (query_id, distances_tra[0])
            assert query_key not in self.traj_collection[key], "Error, exists %s." % query_key
            self.traj_collection[key][query_key] = {
                "state": states_tra,
                "distance": distances_tra,
                "target": tgt_ids,
                "concepts": concepts,
                "rank": rank
            }
            self.c_recorded_sample += 1

    def summary(self):
        print("————————————————————————————————————————————————————————————————————————————————")
        print("Max distance:", self.maximum_distance)

        max_step = self.infer_steps + 3  # one for start, one for failed, one for check
        c_sample_per_distance = {i: 0 for i in range(0, self.maximum_distance + 3)}
        c_sample_per_distance[int(MAX_MOMENT_DISTANCE)] = 0
        distances = list(range(self.minimum_distance, self.maximum_distance + 1))
        if self.opts.mode in ["infer_hero", "infer_conquer"]:
            distances += [int(MAX_MOMENT_DISTANCE)]

        c_ttl = 0
        print("Distance:", self.step_used.keys())
        print("#Step     " + " ".join(["%6d" % i for i in range(0, max_step)]))
        for init_d in distances:
            print("D=%-6d" % init_d, end="  |")
            for num_step_used in range(0, max_step):
                step_c = self.step_used[init_d][num_step_used] if ((init_d in self.step_used) and (num_step_used in self.step_used[init_d])) else 0
                c_sample_per_distance[init_d] += step_c
                print("%4d" % step_c, end=" ")
            print("(%4d)" % c_sample_per_distance[init_d])
            c_ttl += c_sample_per_distance[init_d]
        print("%d samples in total." % c_ttl)
        print()

        print("#Step     " + " ".join(["%6d" % i for i in range(0, max_step)]))
        for init_d in distances:
            print("D=%-6d" % init_d, end="  |")
            for num_step_used in range(0, max_step):
                step_c = self.step_used[init_d][num_step_used] if ((init_d in self.step_used) and (num_step_used in self.step_used[init_d])) else 0
                rate = step_c / c_sample_per_distance[init_d] if c_sample_per_distance[init_d] else 0
                print("%.4f" % rate, end=" ")
            print()
        print()

        ttl = self.c_out_step_nk + self.c_covered_step_nk
        print("Recorded %d samples in total" % self.c_recorded_sample)
        if ttl:
            print("Covered step: %d/%d (%.2f), Out: %d/%d (%.2f)" % (self.c_covered_step_nk, ttl, self.c_covered_step_nk / ttl, self.c_out_step_nk, ttl, self.c_out_step_nk / ttl))

        if self.is_inference:
            action_tag = "N%d" % self.opts.k_for_nk if self.opts.action_type == "nk" else "N1"
            f_path = 'traj/traj_%s_%s_%s.json' % (self.opts.mode, action_tag, self.training_prefix)
            if not os.path.exists("./traj"):
                os.makedirs("./traj")
            with open(f_path, "w") as f:
                json.dump(self.traj_collection, f)
            print('%s saved.' % f_path)


# 파이썬 리스트
attendList = ['순철', '병헌', '민우', '찬호', '민태']
print(attendList)

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print(numbers)

complex = [1, 2.0, 3.1415, 50, '5', "6"]
print(complex)

for data in numbers: # iterable
    print(data)

for data in complex:
    print(complex)

len(numbers)
len(complex)

msg = 'Hello, World!'
print(len(msg))

# 메세지 입력받고 문자열 길이 출력
msg = input('메세지를 입력하세요 ')
print(f'입력한 문자열의 길이: {len(msg)}')

print(len(['Hello, Python!']))
print(len('Hello, Python!'))

# 리스트의 모든 항목 조회
print(complex[0])
print(complex[1])
print(complex[2])
print(complex[3])
print(complex[4])
print(complex[5])

for i in range(6):
    print(complex[i])

for item in complex:
    print(item)

for idx, item in enumerate(complex):
    print(f'{idx}, {item}')

print(complex.index(3.1415))

sports = ['baseball', 'basketball', 'tennis', 'golf', 'soccer']
print(sports.index('soccer'))
print(sports[len(sports) - 1])

languages = ['c/c++', 'c#', 'python', 'java']
print(languages.index('python'))

hobby = ['독서', '등산', '요리']
hobby.append('배구')

print(hobby)

number = [1, 2, 3, 4, 5, 6, 7, 8, 9]
number.append(10)
number.insert(5, 6)
print(number)

weather = ['The', 'weather', 'very']
weather.insert(2, 'is')
weather.insert(4, 'cold')

# 성적 처리 프로그램 - 3명의 학생 데이터를 입력한 뒤,
# 총점, 평균, 학점 처리 후 결과 출력
names = []
kors = []
engs = []
mats = []

for i in range(3):
    name = input('이름은? ')
    names.append(name)
    kor = int(input('국어는? '))
    kors.append(kor)
    eng = int(input('영어는? '))
    engs.append(eng)
    mat = int(input('수학은? '))
    mats.append(mat)

tots = []
avgs = []
grds = []

for i in range(3):
    tots.append(kors[i] + engs[i] + mats[i])
    avgs.append(tots[i] / 3)
    grds.append('가')
    if avgs[i] >= 90: grds[i] = '수'
    if avgs[i] >= 80: grds[i] = '우'
    if avgs[i] >= 70: grds[i] = '미'
    if avgs[i] >= 60: grds[i] = '양'

for i in range(3):
    print(f'{names[i]}, {kors[i]}, {engs[i]}, {mats[i]} \n'
          f'{tots[i]}, {avgs[i]}, {grds[i]} \n')

# 리스트 수정
# 리스트[인덱스] = 수정할 값
hobby
hobby[1] = '여행'

# 리스트 삭제
# pop: 맨 뒤의 항목을 제거
# pop(위치): 해당 위치의 항목을 제거
hobby
hobby.pop()

sports
sports.pop(2)

# remove: 항목값으로 제거
languages
languages.remove('java')

# 과일 리스트에서 야채 삭제하기
fruits = ['사과', '망고', '당근', '수박', '포도', '참외', '토마토']

# 위치값으로 삭제
fruits.pop(2)
fruits.pop(5)
fruits

# 값으로 삭제
fruits.remove('당근')
fruits.remove('토마토')
fruits

# 공인중개사 합격 여부 알아보기
exam = [55, 35, 40, 70, 65, 30]

cnt = len(exam)
sum = 0
fails = 0
result = '아쉽습니다. 다음에 다시 도전해주세요.'

for i in range(cnt):
    if exam[i] < 40: fails += 1 # 과락 수 증가
    sum += int(exam[i])

avg = sum / cnt
if fails == 0 and avg >= 60:
    result = '축하합니다. 합격하셨습니다.'

print(f'평균 점수: {avg:.2f}')
print(result)

# 정렬하기
numbers = [5, 1, 3, 4, 2, 6]

numbers.sort()
numbers

numbers.sort(reverse=True)
numbers

# 모의고사 높은 성적순으로 정렬하기
jumsu = [90, 100, 88, 85, 95, 92, 70, 75, 100, 92, 78, 80, 75, 95, 90, 100, 84]
jumsu.sort(reverse=True)
jumsu

# 문자 정렬 (한글)
names = ['김길동', '이길동', '박길동', '홍길동', '장길동']
names.sort(reverse=True)
names

# 문자 정렬 (영어)
units = ['scv', 'marine', 'firebat', 'ghost', 'dropship', 'battlecruiser', 'valkyrie', 'medic']
units.sort(reverse=True)
units

# 리스트 슬라이싱
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
alphabet.sort(reverse=True)
alphabet

alphabet[2:6] # 2~5까지 추출

alphabet[0:5] # 0~4까지 추출
alphabet[:5] # 0~4까지 추출

alphabet[3:8] # 3~7까지 추출

alphabet[5:] # 5~끝까지 추출
alphabet[5:10] # 5~끝까지 추출

alphabet[3:9] # 3~8까지 추출

alphabet[6:] # 6~끝까지 추출
alphabet[-4:] # 오른쪽을 기준으로 왼쪽 방향 4번째 요소에서 시작

# a b c d e
# 1 2 3 4 5
# -5 -4 -3 -2 -1

# a, b, c, d 추출
alphabet[0:4]
alphabet[:4]
alphabet[:-6]

# d, c, b, a 추출
alphabet[3::-1]
alphabet[-7:-11:-1]
alphabet[-7::-1]

alphabet[-7:-10]
alphabet

# g, h, e, d 추출


import json
import os
import re
import sys

from nltk import casual_tokenize

from src.preprocessor.preprocess import CODE_TOKENISATION_REGEX
from src.tagger.config import Configuration
from src.tagger.data_utils import CorpusIterator
from src.tagger.model import CodePoSModel

def interactive_shell(model, casual=False):
    """Creates interactive shell to play with model
    Args:
        model: instance of NERModel
        casual: If we should use the nltk casual tokenize
    """
    model.logger.info("""
        This is an interactive mode.
        To exit, enter 'exit'.
        You can enter a sentence like
        input> If you have a java.io.InputStream object, how should you process that object and produce a String?""")

    while True:
        sentence = input("input> ")

        if casual:
            words_raw = casual_tokenize(sentence.strip())
        else:
            words_raw = [l.strip()
                         for l in re.findall(CODE_TOKENISATION_REGEX,
                                             sentence.strip())
                         if len(l.strip()) > 0]

        if words_raw == ["exit"]:
            break

        preds = model.predict(words_raw)
        if isinstance(preds, tuple):
            preds = preds[0]

        print(' '.join(['%s_%s' % (w, t) for w, t in zip(words_raw, preds)]))

def restore_model(config):
    try:
        with open(os.path.join(os.path.dirname(config.dir_model), 'config.json')) as f:
            config.__dict__.update(**json.loads(f.read()))
    except FileNotFoundError:
        pass  # We then need the config file to be correct from the start.

    # build model
    model = CodePoSModel(config)
    model.build()
    model.restore_session(config.dir_model)
    return model

def main():
    # create instance of config
    config = Configuration()
    config.dir_model = sys.argv[1]
    model = restore_model(config)

    # create dataset
    test = CorpusIterator(config.filename_test, config.processing_word,
                          config.processing_tag, config.with_l_id, config.max_iter)

    # evaluate and interact
    model.evaluate(test)
    interactive_shell(model)

if __name__ == "__main__":
    main()


#Developed by Henrique Treza

#Venda de produtos com Descontos

preco = float(input('Qual o valor do produto? R$'))
desconto = preco - (preco * 5/100)

print('O produto que tinha o valor de R${:.2f} e '
      'na promoção de 5% de desconto, você irá pagar somente R${}'.format(preco, desconto))

desconto10 = preco - (preco * 10/100)
print('O produto que tinha o valor de R${:.2f} e '
      'na promoção de 10% de desconto, você irá pagar somente R${}'.format(preco, desconto10))

desconto47 = preco - (preco * 47/100)
print('O produto que tinha o valor de R${:.2f} e '
      'na promoção de 47% de desconto, você irá pagar somente R${}'.format(preco, desconto47))

desconto73 = preco - (preco * 73/100)
print('O produto que tinha o valor de R${:.2f} e '
      'na promoção de 73% de desconto, você irá pagar somente R${}'.format(preco, desconto73))


# 끝이 0인 경우는 1,0 붙일수있음, 1로 끝나는경우 0만 추가가능

dp=[[0] * 2 for _ in range(100)]
dp[1]=[0,1]
dp[2]=[1,0]
dp[3]=[1,1]

N=int(input())
for i in range(4,N+1):
    dp[i][0]=dp[i-1][0]+dp[i-1][1] # 0으로 끝나려면 이전꺼의 1로끝나는경우,0으로끝나는경우
    dp[i][1]=dp[i-1][0] # 1로 끝나려면 이전꺼의 0으로 끝나는경우

print(sum(dp[N]))


import asyncio
from functools import wraps
from pathlib import Path

import typer
from rich import print
from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn

from sesg_cli.config import Config
from sesg_cli.database.connection import Session
from sesg_cli.database.models import (
    Experiment,
    SearchStringPerformance,
)

class AsyncTyper(typer.Typer):
    def async_command(self, *args, **kwargs):
        def decorator(async_func):
            @wraps(async_func)
            def sync_func(*_args, **_kwargs):
                return asyncio.run(async_func(*_args, **_kwargs))

            self.command(*args, **kwargs)(sync_func)
            return async_func

        return decorator

app = AsyncTyper(rich_markup_mode="markdown", help="Perform Scopus searches.")

@app.async_command()
async def search(
    experiment_name: str = typer.Argument(
        ...,
        help="Name of the experiment to retrieve the search strings from.",
    ),
    config_file_path: Path = typer.Option(
        Path.cwd() / "config.toml",
        "--config-file-path",
        "-c",
        help="Path to the `config.toml` file.",
    ),
):
    """Searches the strings of the experiment on Scopus."""
    from sesg.evaluation import EvaluationFactory, Study
    from sesg.scopus import InvalidStringError, Page, ScopusClient

    with Session() as session:
        experiment = Experiment.get_by_name(experiment_name, session)
        slr = experiment.slr

        config = Config.from_toml(config_file_path)

        print("Retrieving experiment search strings...")
        search_strings_list = experiment.get_search_strings_without_performance(session)

        evaluation_gs = [
            Study(
                id=s.id,
                title=s.title,
                references=[Study(id=ref.id, title=ref.title) for ref in s.references],
            )
            for s in slr.gs
        ]

        evaluation_qgs = [
            Study(
                id=s.id,
                title=s.title,
                references=[Study(id=ref.id, title=ref.title) for ref in s.references],
            )
            for s in experiment.qgs
        ]

        evaluation_factory = EvaluationFactory(
            gs=evaluation_gs,
            qgs=evaluation_qgs,
        )

        client = ScopusClient(config.scopus_api_keys)

        with Progress(
            TextColumn(
                "[progress.description]{task.description}: {task.completed} of {task.total}"  # noqa: E501
            ),
            BarColumn(),
            TaskProgressColumn(),
        ) as progress:
            overall_task = progress.add_task(
                "Overall",
                total=len(search_strings_list),
            )

            for search_string in search_strings_list:
                progress_task = progress.add_task(
                    "Paginating",
                )

                results: list[Page.Entry] = []

                try:
                    async for page in client.search(search_string.string):
                        progress.update(
                            progress_task,
                            total=page.n_pages,
                            advance=1,
                        )

                        results.extend(page.entries)

                    evaluation = evaluation_factory.evaluate([r.title for r in results])
                    performance = SearchStringPerformance.from_studies_lists(
                        n_scopus_results=len(results),
                        qgs_in_scopus=[
                            slr.get_study_by_id(s.id) for s in evaluation.qgs_in_scopus
                        ],
                        gs_in_scopus=[
                            slr.get_study_by_id(s.id) for s in evaluation.gs_in_scopus
                        ],
                        gs_in_bsb=[
                            slr.get_study_by_id(s.id) for s in evaluation.gs_in_bsb
                        ],
                        gs_in_sb=[
                            slr.get_study_by_id(s.id) for s in evaluation.gs_in_sb
                        ],
                        start_set_precision=evaluation.start_set_precision,
                        start_set_recall=evaluation.start_set_recall,
                        start_set_f1_score=evaluation.start_set_f1_score,
                        bsb_recall=evaluation.bsb_recall,
                        sb_recall=evaluation.sb_recall,
                        search_string_id=search_string.id,
                    )

                    session.add(performance)
                    session.commit()

                except InvalidStringError:
                    print("The following string raised an InvalidStringError")
                    print(search_string.string)

                    performance = SearchStringPerformance(
                        n_scopus_results=-1,
                        qgs_in_scopus=[],
                        gs_in_bsb=[],
                        gs_in_sb=[],
                        n_gs_in_scopus=0,
                        n_qgs_in_scopus=0,
                        gs_in_scopus=[],
                        n_gs_in_bsb=0,
                        n_gs_in_sb=0,
                        start_set_precision=0,
                        start_set_recall=0,
                        start_set_f1_score=0,
                        bsb_recall=0,
                        sb_recall=0,
                        search_string_id=search_string.id,
                    )

                    session.add(performance)
                    session.commit()

                finally:
                    progress.remove_task(progress_task)
                    progress.advance(overall_task)

            progress.remove_task(overall_task)


import re
import logging

from flask import Flask, session, redirect, url_for, escape, request

from .core import commands

app = Flask(__name__)
logger = logging.getLogger(__name__)

@app.route('/')
def index():

    q = request.args.get('q', '')
    if commands.is_url(q):
        # we pull out q this way because we are tricksy here, because q can contain urls
        # and those urls can have ? and & we want to respect that without the person submitting
        # having to worry about encoding the url. So everything after the q is considered part of the
        # q regardless of it has & in it, this also means that we violate the http spec because
        # query string parameters actually have order, shrug
        m = re.search("q=(.*$)", request.environ.get("QUERY_STRING", ""))
        q = commands.unquote(m.group(1))

    v = commands.find(q)
    if commands.is_url(v):
        body = [
            "<!DOCTYPE html>",
            "",
            "    ",
            '        <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />',
            '        <meta http-equiv="refresh" content="0;url=\'{}\'">'.format(v),
            '        <meta name="referrer" content="none">',
            "    ",
            "    ",
            "",
        ]

        logger.info("Redirecting to {}".format(v))

    elif isinstance(v, dict):
        body = [
            "<!DOCTYPE html>",
            "",
            "    ",
            '        <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />',
            '        <meta name="referrer" content="none">',
            "    ",
            '    ',
        ]

        body.append(
            '<form name="autoform" action="{}" method="{}">'.format(v["action"], v.get("method", "POST"))
        )

        for vk, vv in v.items():
            body.append('    <input type="hidden" name="{}" value="{}" />'.format(vk, vv))

        body.append('</form>')
        body.extend(["", "",])

    else:
        body = [
            "<!DOCTYPE html>",
            "",
            "    ",
            '        <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />',
            "        <title>{}</title>".format(q),
            "    ",
            "    ",
            v,
            "    ",
            "",
        ]

    return "\n".join(body)

@app.route("/robots.txt")
def robots():
    #<link rel="search" type="application/opensearchdescription+xml" title="search" href="/opensearch.xml">
    headers = {"Content-Type": "text/plain; charset=utf-8"}
    return ("\n".join([
        "User-agent: *",
        "Disallow: /",
    ]), 200, headers)

@app.route("/opensearch.xml")
def opensearch():
    #<link rel="search" type="application/opensearchdescription+xml" title="search" href="/opensearch.xml">
    headers = {"Content-Type": "text/xml; charset=utf-8"}
    host = url_for("index", _external=True)
    return ("\n".join([
        '<?xml version="1.0" encoding="UTF-8"?>',
        "",
        '<OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">',
        "    <ShortName>bounce</ShortName>",
        "    <Description>Bounce Search</Description>",
        "    <InputEncoding>UTF-8</InputEncoding>",
        #'    <Image height="16" width="16" type="image/x-icon">http://marcyes.com/favicon.ico</Image>',
        "",
        #"    <Contact>name@email.com</Contact>",
        '    <Url type="text/html" template="{}" method="GET">'.format(host),
        '        <Param name="q" value="{searchTerms}"/>',
        #'        <Param name="collection" value="default_collection"/>',
        '        <Param name="opensearch" value="1"/>',
        "    </Url>",
        "    <moz:SearchForm>{}</moz:SearchForm>".format(host),
        "    <SyndicationRight>limited</SyndicationRight>",
        "",
        "    <AdultContent>false</AdultContent>",
        "    <Language>en-us</Language>",
        "</OpenSearchDescription>",
    ]), 200, headers)

@app.route("/list/")
def ls():
    q = request.args.get('q', '')
    if q: q = q.lower()

    lines = []
    lines.append('
')
    lines.append('
')
    lines.append('Keywords	')
    lines.append('domain	')
    lines.append('note	')
    lines.append('
')

    # for display, sort the commands in alphabetical order
    cs = sorted(commands, key=lambda cmd: cmd[0])

    for cmd, val, note in cs:
        if q and q not in cmd.lower(): continue
        lines.append('
')
        lines.append('{}	'.format(cmd))
        lines.append('')
        if re.match("\S://", val):
            lines.append(val)

        else:
            lines.append('
{}
'.format(val))

        lines.append('	')
        lines.append('{}	'.format(note))
        lines.append('
')

    lines.append('
')
    return "\n".join(lines)