from sklearn import mixture

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler
import time
from collections import Counter
import copy

import matplotlib_utils as mu
import pandas_utils as pu
import dimensionality_reduction_utils as dr

from os import sys
sys.path.append("/notebooks/Neurosignal_Final/PRML/")
sys.path.append("/neuron_mesh_tools/Neurosignal_Final/PRML/")
#from prml.rv import VariationalGaussianMixture

# =------------- Functions to Help with Plotting -------------- #

def plot_BIC_and_Likelihood(gmm_data,
                       fig_width=12,
                       fig_height=5,
                           title_suffix = ""):
    """
    Pupose

    """
    cluster_values = list(gmm_data.keys())
    cluster_labels = "Number of Clusters"

    fig1,_ = plt.subplots(1,2)

    fig1 = mu.plot_graph(
        title = "Average Log Likelihood Of Data vs. Number of Clusters\n" + title_suffix,
        y_values = [gmm_data[K]["log_likelihood"] for K in cluster_values],
        x_values = cluster_values,
        x_axis_label = cluster_labels,
        y_axis_label = "Average Log Likelihood of Data (Per Sample)",
        figure = fig1,
        ax_index = 0,
    )

    fig1 = mu.plot_graph(
        title = "BIC vs. Number of Clusters\n" + title_suffix,
        y_values = [gmm_data[K]["bic_value"] for K in cluster_values],
        x_values = cluster_values,
        x_axis_label = cluster_labels,
        y_axis_label = "BIC (Bayesian Information Criterion)",
        figure = fig1,
        ax_index = 1
    )

    fig1.set_tight_layout(True)
    fig1.set_size_inches(fig_width, fig_height)
    return fig1

# ----------------- Functions for analysis ------------------- #
def gmm_analysis(X_train,
                possible_K = list(range(2,8)),
                 reg_covar = 0.00001,
                init_params = "kmeans",
                covariance_type = "full",
                 pca_obj = None,
                 scaler_obj = None,
                 column_titles = None,
                 model_type="mixture", #other type is "variational"
                 verbose=True
                ):

    """
    Purpose: Will perform gmm analysis for a specified different
    number of clusters and save the models and relevant data for further analysis

    """
    if column_titles is None:
        columns_picked = X_train.columns
    else:
        columns_picked = column_titles

    if type(X_train) == pd.DataFrame:
        X_train = X_train.to_numpy()

    scaled_original_results = dict()

    for K in possible_K:
        if verbose:
            print(f"\n\n------Working on clusters K={K}-----")
        st_time = time.time()
        scaled_original_results[K] = dict()

        reg_covar_local = copy.copy(reg_covar)
        #1) Training the GMM
        while reg_covar_local <= 0.1:
            try:
                if model_type == "mixture":
                    if verbose:
                        print("Using mixture model")
                    gmm = mixture.GaussianMixture(n_components=K,
                                                  covariance_type=covariance_type,
                                                 reg_covar=reg_covar,
                                                 init_params=init_params)
                elif model_type == "variational":
                    if verbose:
                        print("Using variational model")
                    raise Exception("Currently not implemented")
                    #gmm = VariationalGaussianMixture(n_components=K)
                else:
                    print("Not right model")
                    raise Exception(f"The gmm model was not picked as mixture or variational : {model_type}")

                gmm.fit(X_train)

            except Exception as e:
                print(f"Exception occured = {str(e)}")
                print(f"Errored on gmm for reg_cov = {reg_covar_local}")
                reg_covar_local = reg_covar_local*10
            else:
                break

        if reg_covar_local >= 1:
            raise Exception(f"No gmm converged and reg_cov was {reg_covar_local}")

        if model_type == "mixture":
            bic_value = gmm.bic(X_train)
            average_log_likelihood_train = gmm.score(X_train)
            current_means = gmm.means_
        else:
            bic_value = 0
            average_log_likelihood_train = 0
            current_means = gmm.mu

        # Getting the Average Log likelihood:

        scaled_original_results[K]["model"] = gmm
        scaled_original_results[K]["log_likelihood"] = average_log_likelihood_train
        scaled_original_results[K]["bic_value"] = bic_value
        scaled_original_results[K]["reg_covar"] = reg_covar

        if not pca_obj is None:
            if verbose:
                print("reversing the pca transformation")
            current_means = pca_obj.inverse_transform(current_means)

        if not scaler_obj is None:
            if verbose:
                print("reversing the normalizing transformation")
            current_means = scaler_obj.inverse_transform(current_means)

        recovered_means = pd.DataFrame(current_means)
        recovered_means.columns = columns_picked

        scaled_original_results[K]["recovered_means"] = recovered_means

        if verbose:
            if model_type == "mixture":
                print(f"Convergence status = {gmm.converged_}")
            print(f"Total time for GMM = {time.time() - st_time}")

    return scaled_original_results

def gmm_classification(gmm_model,curr_data,
                       classification="hard",
                       verbose=True,
                       return_counts=True,
                      ):
    """
    Purpose: Will use the gaussian model passed to
    classify the data points as to which
    cluster they belong

    """
    if type(gmm_model) == mixture.GaussianMixture:
        probs = gmm_model.predict_proba(curr_data)
    elif type(gmm_model) == VariationalGaussianMixture:
        raise Exception("Currently not implemented")
        #probs = gmm_model.classify_proba(curr_data.to_numpy())
    else:
        raise Exception(f"The gmm model was not a mixture or Variational model: {type(gmm_model)}")

    if classification == "soft":
        count_values = np.sum(probs,axis=0)
    elif classification == "hard":
        gmm_class = np.argmax(probs,axis=1)
        counter_obj = Counter(gmm_class)
        count_values = []
        for clust_idx in range(gmm_model.n_components):
            if clust_idx in counter_obj.keys():
                count_values.append(counter_obj[clust_idx])
            else:
                count_values.append(0)
        count_values = np.array(count_values)
    if verbose:
        sorted_cluster_values = np.flip(np.argsort(count_values))
        print(f"Classification: {dict([(k,np.round(count_values[k],2)) for k in sorted_cluster_values])}")

    return count_values

def category_classifications(model,labeled_data,
                                       return_dataframe=True,
                                       verbose = False,
                                       classification_types = ["hard","soft"]):
    total_hard = []
    total_soft = []

    labeled_data_classification = dict()
    dicts_for_classif_df = []

    for c_type in classification_types:

        if verbose:
            print(f"\nclassification_type={c_type}")
        labeled_data_classification[c_type]=dict()

        for k,v in labeled_data.items():

            if verbose:
                print(f"{k}")

            curr_class = gmm_classification(model,v,classification=c_type,verbose=verbose)
            labeled_data_classification[c_type][k] = curr_class

            classifier_dict = dict()
            classifier_dict["classification"]=c_type
            classifier_dict["category"]=k
            classifier_dict["n_clusters"] = model.n_components
            classif_dict_up = dict([(f"cl_{i}",np.round(bb,1)) for i,bb in enumerate(curr_class)])
            classifier_dict.update(classif_dict_up)

            dicts_for_classif_df.append(classifier_dict)

    if return_dataframe:
        # Print out the classification Numbers in Easy to See Dataframe
        df_class = pd.DataFrame.from_dict(dicts_for_classif_df)
        df_class.style.set_caption(f"Clustering Numbers By Neuroscience Category for K = {model.n_components}")
        df_class = df_class.sort_values(by=['category'])
        #print(df_class.to_markdown())

        return labeled_data_classification,df_class
    else:
        return labeled_data_classification

def clustering_stats(data,clust_perc=0.80):
    """
    Will computer different statistics about the clusters
    formed that will be later shown or plotting

    Metrics: For each category and classification type
    1) highest_cluster identify
    2) highest_cluster_percentage
    3) n clusters needed to encompass clust_perc % of the category
    4) Purity statistic

    """
    # categories = ["Apical","Basal","Axon"]
    # classifications = ["hard","soft"]
    # clust_perc = 0.8

    classifications = list(data.keys())
    categories = list(data[classifications[0]].keys())

    stats_dict_by_classification = dict()

    for curr_classification in classifications:
        stats_dict = dict()

        total_per_cluster_by_category = [data[curr_classification][c] for c in categories]
        total_per_cluster = np.sum(total_per_cluster_by_category,axis=0)
        for curr_category in categories:
            local_stats_dict = dict()

            count_data = data[curr_classification][curr_category]

            """
            Statistics to find:
            1) The cluster with the most of that label and the % in that cluster
            2) The number of clusters needed to comprise 80% of labeled group
            3) The purity measurements

            Pseudocode:
            1) get the total number items put in each cluster across all categories
            2) For each cluster:
            a. Multiply the perc in that cluster * (curent number in that cluster/total number in that cluster)

            """

            sorted_labels = np.flip(np.argsort(count_data))
            highest_cluster_perc = count_data[sorted_labels[0]]/np.sum(count_data)

            local_stats_dict["highest_cluster"]  = sorted_labels[0]
            local_stats_dict["highest_cluster_perc"] = highest_cluster_perc

            sorted_labels_cumsum_perc = np.cumsum(count_data[sorted_labels]/np.sum(count_data))
            perc_per_cluster = count_data/np.sum(count_data)

            n_clusters = np.digitize(clust_perc,sorted_labels_cumsum_perc)+1
            local_stats_dict[f"n_clusters_{np.floor(clust_perc*100)}"] = n_clusters

            #find the purity metric

            purity = np.sum(perc_per_cluster[total_per_cluster != 0]*count_data[total_per_cluster != 0]/total_per_cluster[total_per_cluster != 0])

            local_stats_dict["purity"] = purity

            stats_dict[curr_category] = local_stats_dict

        # measure the purity of each cluster
        max_per_cluster = cluster_purity = np.max(total_per_cluster_by_category,axis=0)
        cluster_purity = [m/t_c if t_c > 0 else 0 for m,t_c in zip(max_per_cluster,total_per_cluster)]

        stats_dict_by_classification[curr_classification] = dict(cluster_purity= cluster_purity,stats_dict=stats_dict)
    return stats_dict_by_classification

def cluster_stats_dataframe(labeled_data_classification):
    """
    Purpose: Just want to visualize the soft and the hard assignment (and show they are not that different)

    Pseudocode:
    1)

    """

    ret_stats = clustering_stats(labeled_data_classification)

    dict_for_df = []

    for cl_type,cl_data in ret_stats.items():
        k = len(cl_data["cluster_purity"])
        curr_stats_dict = cl_data["stats_dict"]

        for cat_name,cat_stats_dict in curr_stats_dict.items():
            cat_local_dict = dict()
            cat_local_dict["category"] = cat_name
            cat_local_dict["classification"] = cl_type
            cat_local_dict["n_clusters"] = k
            cat_local_dict.update(cat_stats_dict)

            dict_for_df.append(cat_local_dict)

    df = pd.DataFrame.from_dict(dict_for_df)
    return df.sort_values(by=['category'])

def plot_advanced_stats_per_k(advanced_stats_per_k,
                             stats_to_plot = ["highest_cluster_perc","purity"],
                              title_suffix="",
                             fig_width = 12,
                              fig_height = 5):
    """
    Purpose: plotting the highest cluster and purity as a function of k

    Pseudocode:
    0) Get all the possible categories, n_clusters
    0) Sort by n_clusters
    1) Iterate through all the stats we want to plot
        2) Iterate through all of the categories
            -- for all n_clusters
            a. Restrict by category and n_clusters and pull down the statistic
            b. Add to list
            --
            c. Save full list in dictionary

        3) Plot the stat using the category dictionary (using the ax index id)

    """

    advanced_stats_df = pd.concat(list(advanced_stats_per_k.values()))

    unique_categories = np.unique(advanced_stats_df["category"].to_numpy())
    unique_n_clusters = np.unique(advanced_stats_df["n_clusters"].to_numpy())
    stats_to_plot = ["highest_cluster_perc","purity"]

    cluster_labels = "Number of Clusters"

    fig, _ = plt.subplots(1,len(stats_to_plot))

    for j,st in enumerate(stats_to_plot):
        st_cat_dict = dict()
        for cat in unique_categories:
            cat_list = []
            for k in unique_n_clusters:
                curr_st = advanced_stats_df.query(f"category=='{cat}' & n_clusters=={k}")[st].to_numpy()
                if len(curr_st) != 1:
                    raise Exception("Stat was not of size 1")
                cat_list.append(curr_st[0])
            st_cat_dict[cat] = cat_list

            fig = mu.plot_graph(
                title = f"{st} vs. Number of Clusters\n" + title_suffix,
                y_values = cat_list,
                x_values = unique_n_clusters,
                x_axis_label = cluster_labels,
                y_axis_label = f"{st}",
                figure = fig,
                ax_index = j,
                label=cat
            )
    fig.set_tight_layout(True)
    fig.set_size_inches(fig_width, fig_height)
    return fig

def gmm_pipeline(df,
                title_suffix,
                labeled_data_indices=None,
                 category_column=None,
                 columns_picked=None,
                 possible_K = list(range(2,8)),
                 print_tables = None, #clusters will print the clustering tables fro
                 apply_normalization=True,
                 apply_pca = True,
                 pca_whiten=True,
                 plot_sqrt_eigvals=True,
                 n_components_pca = None,
                 classification_types = ["hard"],#["hard","soft"]
                 model_type = "mixture",
                 verbose=True,

                ):
    """
    Will carry out all of the clustering analysis and
    advanced stats analysis on a given dataset

    Arguments:
    A data table with all of the labeled data
    """
    # ------- Initializing variables --------- #

    if labeled_data_indices is None:
        category_col = df[category_column].to_numpy()
        unique_labels = np.unique(category_col)
        labeled_data_indices = dict([(k,np.where(category_col==k)[0]) for k in unique_labels])
        df = pu.delete_columns(df,[category_column])

    if print_tables is None:
        print_tables = possible_K

    # -------- Part 0: Preprocessing (Column restriction, Normalization, PCA) ----------- #
    if verbose:
        print(f"# -------- Part 0: Preprocessing (Column restriction, Normalization, PCA) ----------- #")

    if columns_picked is None:
        columns_picked = list(df.columns)
    else:
        if verbose:
            print(f"Restricting to columns : {columns_picked}")
        df = df[columns_picked]

    # Scaling the Data
    if apply_normalization:
        if verbose:
            print(f"Applying Normalization")

        scaler_obj = StandardScaler()
        df_data_scaled = scaler_obj.fit_transform(df)

        #df_data_reversed = scaler.inverse_transform(df_data_scaled,copy=True)

        data_df_normalized = pd.DataFrame(df_data_scaled)
        #add on the columns
        data_df_normalized.columns = df.columns
        df = data_df_normalized
    else:
        scaler_obj = None

    # Applying pca to the data
    if apply_pca:

        if n_components_pca is None:
            n_components_pca=len(columns_picked)

        if verbose:
            print(f"Applying pca with {n_components_pca} components")

        data_analyzed = dr.pca_analysis(df.to_numpy(),
                                    n_components=n_components_pca,
                                    whiten=pca_whiten,
                                    plot_sqrt_eigvals=plot_sqrt_eigvals)

        if verbose:
            print(f'Explained Variance = {data_analyzed["percent_variance_explained_up_to_n_comp"]}')
            dr.plot_variance_explained(data_analyzed)

        df_pca = pd.DataFrame(data_analyzed["data_proj"])
        df_pca.columns = [f"PC_{j}" for j in range(n_components_pca)]
        df = df_pca

        pca_obj = data_analyzed["pca_obj"]

    else:
        pca_obj = None

    # -------- Part 1: GMM clustering with different Number of Clusters ----------- #
    if verbose:
        print(f"# -------- Part 1: GMM clustering with different Number of Clusters ----------- # ")

    X_train = df
    scaled_original_results = gmm.gmm_analysis(X_train,
                    scaler_obj=scaler_obj,
                    pca_obj=pca_obj,
                     possible_K = possible_K,
                    column_titles=columns_picked,
                    model_type = model_type)

    if model_type == "mixture":
        fig1 = gmm.plot_BIC_and_Likelihood(scaled_original_results,title_suffix=title_suffix)
        mu.display_figure(fig1)

    # --------- Part 2: computing the advanced statistics on the clustering ------- #
    if verbose:
        print(f"# --------- Part 2: computing the advanced statistics on the clustering ------- # ")

    advanced_stats_per_k = dict()
    labeled_data = dict([(kk,df.iloc[vv]) for kk,vv in labeled_data_indices.items()])

    for curr_K in scaled_original_results.keys():
        if verbose:
            print(f"\n\n----Working on Advanced Statistics for n_clusters = {curr_K}----\n")

        model = scaled_original_results[curr_K]["model"]

        labeled_data_classification,df_class=gmm.category_classifications(
                                    model,
                                    labeled_data,
                                    classification_types=classification_types)
        if curr_K in print_tables:
            print("Recovered Means From Clustering")
            pu.display(scaled_original_results[curr_K]["recovered_means"])
            print("\n")
            print(f"Clustering Numbers By Neuroscience Category for K = {model.n_components}")
            pu.display_df(df_class)
            print("\n")

        cl_stats_df = gmm.cluster_stats_dataframe(labeled_data_classification)

        if curr_K in print_tables:
            print(f"Clustering Advanced Statistics By Neuroscience Category for K = {curr_K}")
            pu.display_df(cl_stats_df)

        column_restriction = ["category","highest_cluster_perc","purity","n_clusters"]
        cl_stats_restricted = cl_stats_df.query("classification=='hard'")[column_restriction]

        advanced_stats_per_k[curr_K] = cl_stats_restricted

    # -------- Part 3: Plotting the Advanced Cluster Statistics -------------- #
    if verbose:
        print(f"# -------- Part 3: Plotting the Advanced Cluster Statistics -------------- # ")
    fig_current = gmm.plot_advanced_stats_per_k(advanced_stats_per_k,title_suffix=title_suffix)
    mu.display_figure(fig_current)

import gmm


"""
Defines module steppables

Steppables
==========

NeighborRecoverySteppable
-------------------------
Description: implements recovery model

Usage: In ViralInfectionVTM.py, add the following

from Models.RecoveryNeighbor.RecoverySteppables import NeighborRecoverySteppable

CompuCellSetup.register_steppable(steppable=NeighborRecoverySteppable(frequency=1))

NeighborRecoveryDataSteppable
-----------------------------
Description: performs data tracking

Usage: In ViralInfectionVTM.py, add the following

from Models.RecoveryNeighbor.RecoverySteppables import NeighborRecoveryDataSteppable

CompuCellSetup.register_steppable(steppable=NeighborRecoveryDataSteppable(frequency=1))
"""

import random

from . import RecoveryInputs
rec_steppable_key = "nbrec_steppable"
rec_data_steppable_key = "nbrec_data_steppable"

# Inherit from Simple Recovery model
from Models.RecoverySimple.RecoverySteppables import SimpleRecoverySteppable, SimpleRecoveryDataSteppable

class NeighborRecoverySteppable(SimpleRecoverySteppable):
    """
    Implements neighbor-dependent recovery
    """

    unique_key = rec_steppable_key

    def __init__(self, frequency=1):
        super().__init__(frequency)

        self._voxel_area = self.voxel_length * self.voxel_length

        # Initialize default data
        from ViralInfectionVTMSteppables import uninfected_type_name, dead_type_name
        self.register_recovery(dead_type_name=dead_type_name,
                               recovered_type_name=uninfected_type_name,
                               recovery_rate=RecoveryInputs.recovery_rate)

    def start(self):
        super().start()

        # Reducing a little computational cost here
        self._voxel_area = self.voxel_length * self.voxel_length

    def cell_recovers(self, _cell) -> bool:
        """
        Recovery criterion

        :param _cell: dead cell to test for recovery
        :return: True if cell recovers
        """
        recovery_type = self.recovery_type(_cell)
        ca = sum([a for n, a in self.get_cell_neighbor_data_list(_cell) if n is not None and n.type == recovery_type])
        if ca == 0:
            return False
        return random.random() < ca * self._voxel_area * self.recovery_prob(_cell)

class NeighborRecoveryDataSteppable(SimpleRecoveryDataSteppable):
    """
    Implements neighbor-dependent recovery data tracking; like SimDataSteppable in main framework
    """

    unique_keys = rec_data_steppable_key

    def __init__(self, frequency=1, plot_rec_data_freq=None, write_rec_data_freq=None):
        if plot_rec_data_freq is None:
            plot_rec_data_freq = RecoveryInputs.plot_rec_data_freq
        if write_rec_data_freq is None:
            write_rec_data_freq = RecoveryInputs.write_rec_data_freq
        super().__init__(frequency, plot_freq=plot_rec_data_freq, write_freq=write_rec_data_freq)

        # Particularize SimpleRecoveryDataSteppable
        self.rec_steppable_key = rec_steppable_key
        self.data_path_rel = 'nbrec_data.dat'
        self.plot_title = 'Neighbor Recovery'


#----------
# Devin Suy
#----------

import copy
import sys

# Human player is min
# Bot is max
class GameState:
    num_states = -1
    game_states = []

    # turn = True : Successor states will add an X (Human)
    # turn = False : Successor states will add a O (Bot)
    def __init__(self, game_state, turn, depth, is_human=None, cell_num=None):
        self.game_state = copy.deepcopy(game_state)

        # Only called for successor constructor
        if is_human is not None:
            # Add in the change from the parent state
            if is_human:
                player = self.game_state.human
            else:
                player = self.game_state.bot
            self.game_state.board.set_cell(cell_num, player)

        self.turn = turn
        self.depth = depth
        self.successors = []  # Represent "children" states of our current state
        self.state_score = None
        self.leaf_node = False

        if self.has_successors():
            self.generate_succesors()
        else:
            # We have reached the bottom of this subtree
            self.leaf_node = True

        # Save states
        GameState.num_states += 1
        GameState.game_states.append(self)
        print(GameState.num_states)

    def has_successors(self):
        if (
            len(self.game_state.board.avail_cell_nums) > 0
            and (not self.game_state.game_over(self.game_state.human))
            and (not self.game_state.game_over(self.game_state.bot))
        ):
            return True

        return False

    def generate_succesors(self):
        # Create a successor for placement in every possible square for this turn
        for cell_num in self.game_state.board.avail_cell_nums:
            # Recursively generate successor states
            successor_state = copy.deepcopy(self.game_state)
            self.successors.append(GameState(successor_state, not self.turn, self.depth+1, self.turn, cell_num))


#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# @Time    : 2017/11/6 14:42
# @Author  : Bob
# @Website : www.nagexiucai.com
# @E-Mail  : me@nagexiucai.com
# @Summary : 通过简单的狗类演示类的三大性质。

from grasp.classes.horse.simple import Horse

class HorseI(Horse):
    def __init__(self, number):
        self.number = number
        self.__age = "secret"


import warnings
import logging
import bot

chat_commands = {}
default_class = None

class ChatCommand(object):
    '''a command which is run in response to a chat command

    very non-OOP class for representing a chat command
    note: define a sub-class then call SubClass.register(..)
    instances are created through ChatCommand.process_text()
    do not manually create instances of this class or any inheriting
    classes!

    Attributes:
        args: array of arguments for the command executing
        chan: discord.py channel for the context this command executes in
        sender: discord.py user for the context ...
    '''
    # the prefix can be re-defined wherever, possibly per-chanhel-connection?
    command_prefix = "!"
    name = "PLACEHOLDER"

    def __init__(self):
        """setup new instance with placeholder values"""
        self.args = []
        self.chan = None
        self.sender = None

    def bind_args(self, tokens):
        """bind command arguments in an array

        note self.args[0] == registered name (like sys.argv)

        Args:
            tokens:
        """
        for t in tokens:
            self.args.append(t)

    def bind_context(self, channel, sender):
        """bind discord context for execution.

        note: this MUST be called before execute!

        Args:
            channel: discord.py channel object
            sender: ... user object
        """
        self.chan = channel
        self.sender = sender

    def _body(self):
        """override this for commands to implement execution functionality
        """
        warnings.warn("Default body in ChatCommand")
        return

    def execute(self):
        """perform the action associated with this command.

        do not override this in child classes. this is where logic for valid
        state is done because we cannot do that at instance construction!
        """
        if self.chan is None or self.sender is None:
            raise RuntimeError("Invalid ChatCommand state: no context bound")
        self._body()

    def _helptext(self):
        """
        override with explanatory text for a given command. used for generating
        help messages.
        """
        return "PLACEHOLDER"

    @classmethod
    def register(cls, name):
        """ register this CLASS with a given name, to be looked up in parsing

        Args:
            name: string name associating this class to a discord chat command
        """
        global chat_commands
        if name in chat_commands:
            # override for duplicate names, but log it!
            warnings.warn("Duplicate command registered - overriding "+name,
                          stacklevel=3)
        chat_commands[name] = cls
        cls.name = name

    @classmethod
    def process_text(cls, text):
        """parse text and check for a valid registered command call

        Args:
            text: string message from discord

        Returns:
            ChatCommand instance which can be executed, even if the command
                given in text was invalid (see set_default)
        """
        global chat_commands
        global default_class
        # early out for empty strings
        if len(text) < 0:
            return default_class()
        if text[0] == cls.command_prefix:
            # tokenize without prefix character
            tokens = text[1:].split(" ")
            # execute command name if known otherwise log
            if tokens[0] in chat_commands:
                cmd = chat_commands[tokens[0]]()
                cmd.bind_args(tokens)
                if cmd is None:
                    print("Somehow matched none command! {}".format(tokens[0]))
                return cmd
            else:
                logging.info("Unregistered command {} called".format(
                    tokens[0]))
                print("Returning default command handler class")
                return default_class()
        else:
            return default_class()

    @classmethod
    def set_default(cls):
        """set default class for process_text() commands where no command is found
        """
        print("Setting default chat command class globally")
        global default_class
        default_class = cls
        print(default_class)

class CommandNull(ChatCommand):
    """
    no-action "fallthrough" class for commands which do not map to classes
    """
    def _body(self):
        pass

# set this class as the default!
CommandNull.set_default()

class CommandsHelp(ChatCommand):
    """
    command for listing all commands
    """
    def _body(self):
        global chat_commands
        msg = "Commands list: \n```"
        for key in chat_commands:
            # do not print a message for ourself
            if key == self.__class__.name:
                continue
            # get helptext from temp instance of this class
            c = chat_commands[key]()
            msg += "{}: {}\n".format(key, c._helptext())
        # finally send message using bot.client (logged in client)
        bot.client.message_send(self.chan, msg)


print("give a maximum number ")
n = int(input())
def isPrime(number):
    for i in range(2,int(number/2+1)):
        if number % i == 0:
            return False
    return True

for i in range(2,n+1):
    if isPrime(i) == True:
        print(str(i), end=', ')


# wapf to find fact of an integer

def fact(num):
	f = 1
	for i in range(1, num+1):
		f = f * i
	return f

n = 12
r = 2

perm = fact(n) / fact(n-r)
comb = fact(n) / (fact(r) * fact(n-r))

print("perm = ", perm)
print("comb = ", comb)

# dev karo ek baar.. call karo baar baar
# DRY ==> dont repeat yourself


import helpers
import minkit
import os
import pytest

@pytest.mark.core
def test_backend():
    '''
    Test the construction of the backend.
    '''
    with pytest.raises(AttributeError):
        minkit.Backend.DataSet

    bk = minkit.Backend(minkit.backends.core.CPU)

    x = minkit.Parameter('x', bounds=(-1, +1))

    data = helpers.rndm_gen.uniform(0, 1, 1000)

    # Test initialization and constructor methods
    bk.DataSet(minkit.darray.from_ndarray(data, bk), [x])

    dataset = bk.DataSet.from_ndarray(data, x)

    new_bk = minkit.Backend(minkit.backends.core.CPU)

    m = bk.Parameter('m')
    c = bk.Parameter('c')
    s = bk.Parameter('s')
    k = bk.Parameter('k')
    y = bk.Parameter('y')

    g = bk.Gaussian('gauss', m, c, s)
    e = bk.Exponential('exponential', m, k)

    bk.AddPDFs.two_components('pdf', g, e, y)

    # Test the adaption of objects to new backends
    dataset.to_backend(new_bk)

BACKEND = os.environ.get('MINKIT_BACKEND', None)

if minkit.backends.core.is_gpu_backend(BACKEND):

    def test_gpu_backends():
        '''
        Test the change of objects from a CPU to a GPU backend.
        '''
        minkit.Backend()
        minkit.Backend(BACKEND)


import discord
from discord.ext import commands
import json
from discord.ui import Button,View,Select

with open('setting.json', mode = 'r', encoding = 'utf8') as jfile:
    jdata = json.load(jfile)

class React(commands.Cog):
    def __init__(self,bot :commands.Bot) -> None:
        self.bot = bot

    @commands.command()
    async def button_test(self,ctx):
        view = View()
        button = Button(label='點我啊臭ㄐㄐ',style=discord.ButtonStyle.primary,emoji='<:ji:816467849357951058>')

        async def button_callback(interaction):
            user = interaction.user.name
            await interaction.response.send_message(f'{user}點了，{user}是臭ㄐㄐ')
        button.callback = button_callback
        view.add_item(button)
        await ctx.send(view=view)

    @commands.command()
    async def select_test(self,ctx):
        select = Select(
            placeholder='你覺得蔡神有多帥(滿分100)',
            options=[
                discord.SelectOption(
                    label='100分',
                    description='滿分，蔡神超帥'
                ),
                discord.SelectOption(
                    label='200分',
                    description='蔡神怎麼可能被100分侷限住'
                ),
                discord.SelectOption(
                    label='114514分',
                    description='這麼臭的選項有存在的必要嗎(惱'
                )
            ]
        )
        async def callback(interaction):
            user = interaction.user.name
            await interaction.response.send_message(f'{user}覺得蔡神有{select.values[0]}帥')

        select.callback = callback
        view = View()
        view.add_item(select)
        await ctx.send(view = view)

async def setup(bot):
    await bot.add_cog(React(bot))


import sys
import csv
from PySide2 import QtWidgets as Qtw
from PySide2 import QtGui as Qtg
from PySide2 import QtCore as Qtc

class CSVTableModel(Qtc.QAbstractTableModel):

    def __init__(self, csv_file):
        super(CSVTableModel, self).__init__()

        self.filename = csv_file

        with open(self.filename) as fh:
            csvreader = csv.reader(fh)
            self._headers = next(csvreader)
            self._data = list(csvreader)

if __name__ == '__main__':
    app = Qtw.QApplication(sys.argv)
    mw = CSVTableModel()
    mw.show()
    sys.exit(app.exec_())


import random
import sys
import os

# aceita o número de testes como um comando de parametro linha
tests = int(sys.argv[1])
# aceita os parametros para os teste como um parametro de comando linha
n = int(sys.argv[2])

for i in range(tests):
    print('Test #' + str(i))
    # executa o gerador gen.py com um parametro n e seed i
    os.system('PY Stress2.py ' + str(n) + ' ' + str(i) + ' > input.txt')
    # executa a solução modelo model.py
    os.system('PY max_pairwise_product.py <input.txt >model.txt')
    # executa a solução principal
    os.system('PY max_pairwise_product_fast.py <input.txt >main.txt')

    # lê a saida da solução modelo:
    with open('model.txt') as f:
        model = f.read()
    print('Model: ', model)
    # leia a saida da solução principal:
    with open('main.txt') as f:
        main = f.read()
    print('Main: ', main)
    if model != main:
        break


# -*- coding: utf-8 -*-
"""
Created on Thu May 30 22:00:30 2019

@author: Lawrence
"""
shipping=True
amountPurchase=int(input("Enter your deposit amount:"))
if amountPurchase <50:
    shipping=False
    print("Please add an extra 10$")
print("Shipping is free 0.00$")


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

import env

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler

sql = """
SELECT prop.*,
       pred.logerror,
       pred.transactiondate,
       air.airconditioningdesc,
       arch.architecturalstyledesc,
       build.buildingclassdesc,
       heat.heatingorsystemdesc,
       landuse.propertylandusedesc,
       story.storydesc,
       construct.typeconstructiondesc

FROM   properties_2017 prop
       INNER JOIN (SELECT parcelid,
       					  logerror,
                          Max(transactiondate) transactiondate
                   FROM   predictions_2017
                   GROUP  BY parcelid, logerror) pred
               USING (parcelid)
       LEFT JOIN airconditioningtype air USING (airconditioningtypeid)
       LEFT JOIN architecturalstyletype arch USING (architecturalstyletypeid)
       LEFT JOIN buildingclasstype build USING (buildingclasstypeid)
       LEFT JOIN heatingorsystemtype heat USING (heatingorsystemtypeid)
       LEFT JOIN propertylandusetype landuse USING (propertylandusetypeid)
       LEFT JOIN storytype story USING (storytypeid)
       LEFT JOIN typeconstructiontype construct USING (typeconstructiontypeid)
WHERE  prop.latitude IS NOT NULL
       AND prop.longitude IS NOT NULL AND transactiondate <= '2017-12-31'
"""

def get_db_url(database):
    from env import host, user, password
    url = f'mysql+pymysql://{user}:{password}@{host}/{database}'
    return url

# acquire zillow data using the query
def get_zillow(sql):
    url = get_db_url('zillow')
    zillow_df = pd.read_sql(sql, url, index_col='id')
    return zillow_df

def handle_missing_values(df, prop_required_column = .5, prop_required_row = .70):
	#function that will drop rows or columns based on the percent of values that are missing:\
	#handle_missing_values(df, prop_required_column, prop_required_row
    threshold = int(round(prop_required_column*len(df.index),0))
    df.dropna(axis=1, thresh=threshold, inplace=True)
    threshold = int(round(prop_required_row*len(df.columns),0))
    df.dropna(axis=0, thresh=threshold, inplace=True)
    return df

def remove_columns(df, cols_to_remove):
	#remove columns not needed
    df = df.drop(columns=cols_to_remove)
    return df

def wrangle_zillow():
    df = pd.read_csv('zillow.csv')

    # Restrict df to only properties that meet single unit use criteria
    single_use = [261, 262, 263, 264, 266, 268, 273, 276, 279]
    df = df[df.propertylandusetypeid.isin(single_use)]

    # Restrict df to only those properties with at least 1 bath & bed and 350 sqft area
    df = df[(df.bedroomcnt > 0) & (df.bathroomcnt > 0) & ((df.unitcnt<=1)|df.unitcnt.isnull())\
            & (df.calculatedfinishedsquarefeet>350)]

    # Handle missing values i.e. drop columns and rows based on a threshold
    df = handle_missing_values(df)

    # Add column for counties
    df['county'] = np.where(df.fips == 6037, 'Los_Angeles',
                           np.where(df.fips == 6059, 'Orange',
                                   'Ventura'))
    # drop columns not needed
    df = remove_columns(df, ['id',
       'calculatedbathnbr', 'finishedsquarefeet12', 'fullbathcnt', 'heatingorsystemtypeid'
       ,'propertycountylandusecode', 'propertylandusetypeid','propertyzoningdesc',
        'censustractandblock', 'propertylandusedesc','heatingorsystemdesc','unitcnt'
                            ,'buildingqualitytypeid'])

    # replace nulls with median values for select columns
    df.lotsizesquarefeet.fillna(7313, inplace = True)

    # Columns to look for outliers
    df = df[df.taxvaluedollarcnt < 5_000_000]
    df[df.calculatedfinishedsquarefeet < 8000]

    # Just to be sure we caught all nulls, drop them here
    df = df.dropna()

    return df

def min_max_scaler(train, valid, test):
    '''
    Uses the train & test datasets created by the split_my_data function
    Returns 3 items: mm_scaler, train_scaled_mm, test_scaled_mm
    This is a linear transformation. Values will lie between 0 and 1
    '''
    num_vars = list(train.select_dtypes('number').columns)
    scaler = MinMaxScaler(copy=True, feature_range=(0,1))
    train[num_vars] = scaler.fit_transform(train[num_vars])
    valid[num_vars] = scaler.transform(valid[num_vars])
    test[num_vars] = scaler.transform(test[num_vars])
    return scaler, train, valid, test

def outlier_function(df, cols, k):
	#function to detect and handle oulier using IQR rule
    for col in df[cols]:
        q1 = df.annual_income.quantile(0.25)
        q3 = df.annual_income.quantile(0.75)
        iqr = q3 - q1
        upper_bound =  q3 + k * iqr
        lower_bound =  q1 - k * iqr
        df = df[(df[col] < upper_bound) & (df[col] > lower_bound)]
    return df

def get_mall_customers(sql):
	    url = get_db_url('mall_customers')
	    mall_df = pd.read_sql(sql, url, index_col='customer_id')
	    return mall_df

def wrangle_mall_df():

    # acquire data
    sql = 'select * from customers'

    # acquire data from SQL server
    mall_df = get_mall_customers(sql)

    # handle outliers
    mall_df = outlier_function(mall_df, ['age', 'spending_score', 'annual_income'], 1.5)

    # get dummy for gender column
    dummy_df = pd.get_dummies(mall_df.gender, drop_first=True)
    mall_df = pd.concat([mall_df, dummy_df], axis=1).drop(columns = ['gender'])
    mall_df.rename(columns= {'Male': 'is_male'}, inplace = True)
    # return mall_df
    return mall_df
    # split the data in train, validate and test
    # train, test = train_test_split(mall_df, train_size = 0.8, random_state = 123)
    # train, validate = train_test_split(train, train_size = 0.75, random_state = 123)

    # return min_max_scaler, train, validate, test


def roman_num(func, n: str, value: int = None) -> int:
    def read(n):
        with open("roman.txt", "r") as foo:
            data = foo.read()
            data = eval(data)
        val = data.get(n)
        if val != None:
            print(f"{n} = {val}")
        else:
            print("Value not found")

    def write(n, value):
        with open("roman.txt", "r") as foo:
            data = foo.read()
            if data == "" or data == "{}":
                data = dict()
            else:
                data = eval(data)

        data.update({n: value})
        with open("roman.txt", "w") as foo:
            foo.write(str(data))
        print("Data saved!")

    if func == "read":
        read(n)
    elif func == "write":
        write(n, value)
    else:
        print("Incorrect mode provided!")

mode = input("Enter the function to be carried out (read/write): ")

n = input("Enter Roman numeral: ")

if mode == "write":
    value = input("Enter value: ")
else:
    value = None

roman_num(func = mode, n = n, value = value)


# levanto el archivo
from openpyxl import load_workbook
import func.som as som
import numpy as np

wb = load_workbook('Ejercicios/Act_4/Drug4.xlsx', read_only=True)
ws = wb.get_sheet_by_name('Sheet1')
ws.max_row
ws.max_column
numerical_data = ws.iter_rows(min_row=1, max_row=ws.max_row, min_col=1, max_col=ws.max_column)
muestra = np.array([[cell.value for cell in row] for row in numerical_data])
# separo los datos de la calse
X = muestra[:, 0:6]
T = muestra[:, 6]
# escalamos los datos
X_escalado = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0))
# Genero matriz de valores esperados
T_matriz = np.concatenate(([T == 0], [T == 1], [T == 2], [T == 3]), axis=0).astype(int)

# Defino Parámetros de la red
P = X_escalado.T
columnas = 4
filas = 4
alfaInicial = 0.2
vecindad = 3
func_vecindad = 1  ####otro nro es sombrero
sigma = 2  ##ancho del sombrero
ite_reduce = 50
dibujar = True

matriz = np.zeros((200, 200))
for func_vecindad in [0, 1]:
    for vecindad in [2, 3]:
        for sigma in [1, 2, 3]:
            for columnas in [2, 3, 4]:
                (w_O) = som.train(P, filas, columnas, alfaInicial, vecindad, func_vecindad, sigma, ite_reduce, dibujar)

                # una columna por cada dimension de mi dataset
                ganador = []
                for p in range(200):
                    distancias = -np.sqrt(
                        np.sum((w_O - (P[:, p][np.newaxis]) * np.ones((filas * columnas, 1))) ** 2, 1))
                    ganadora = np.argmax(distancias)
                    ganador.append(ganadora)

                for i in range(200):
                    for j in range(200):
                        ##if i!=j: opcional.....la funcion dendograma se encarga de esto
                        if ganador[i] == ganador[j]:
                            matriz[i][j] += 1

matriz2 = matriz

som.dendograma(matriz, T)


# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://docs.scrapy.org/en/latest/topics/item-pipeline.html

# useful for handling different item types with a single interface
from itemadapter import ItemAdapter
import redis

class PythonscrapyPipeline:

    def process_item(self, item, spider):

        # 建立连接
        pool = redis.ConnectionPool(host='localhost', port=6379, decode_responses=True)
        conn = redis.Redis(connection_pool=pool)

        # 取出标题
        for key, value in item.items():
            value1 = value
            break
        print('目标网站:' + value1)

        for key, value in item.items():
            conn.hset(value1, key, value)

        print(value1 + '爬取成功!')
        # print(conn.hgetall(value1))
        return item


import requests
from flask import Blueprint, render_template, abort, request, session

from src.data_interface import model

bp = Blueprint("you", __name__)

@bp.route("/you")
def index():
    session_id = session["uid"]

    res = requests.get(
        url="http://api:8000/recommend/session/{id}".format(id=session_id),
        params={"page_type": "you"}
    )

    if res.status_code != 200:
        abort(res.status_code)

    res_json = res.json()

    recommendations = res_json["recommendations"]

    return render_template(
        "you/index.html",
        recommendations=recommendations
    )

@bp.route("/you/taste", methods=("GET", "POST"))
def taste():
    if request.method == "POST":
        form_data = request.form.to_dict(flat=False)
        user_genres = form_data.get("genre") or []

        model.User.query\
            .filter_by(username=session.get("username"))\
            .update({"favorite_genres": map(int, user_genres)})

    genres = model.Genre.query.all()
    user_genres = model.User.query\
        .with_entities(model.User.favorite_genres)\
        .filter(model.User.username == session["username"])\
        .one()[0]

    return render_template(
        "you/taste.html",
        genres=genres,
        user_genres=user_genres
    )


from nowstagram import app
from qiniu import Auth,put_data

#需要填写自己的access_key和secret_key
access_key=app.config['QINIU_ACCESS_KEY']
secret_key=app.config['QINIU_SECRET_KEY']

#构建鉴权对象
q=Auth(access_key,secret_key)

#要上传到的空间
bucket_name=app.config['QINIU_BUCKET_NAME']
domain_prefix=app.config['QINIU_DOMAIN']

def qiniu_upload_file(source_file,save_file_name):
    """upload_token生成上传凭证

    Args:
        bucket:  上传的空间名
        key:     上传的文件名，默认为空
        expires: 上传凭证的过期时间，默认为3600s
        policy:  上传策略，默认为空

    Returns:
        上传凭证
    """
    token=q.upload_token(bucket_name,save_file_name)

    """put_data上传二进制流到七牛

    Args:
        up_token:         上传凭证
        key:              上传文件名
        data:             上传二进制流
        params:           自定义变量，规格参考 http://developer.qiniu.com/docs/v6/api/overview/up/response/vars.html#xvar
        mime_type:        上传数据的mimeType
        check_crc:        是否校验crc32
        progress_handler: 上传进度

    Returns:
        一个dict变量，类似 {"hash": "<Hash string>", "key": "<Key string>"}
        一个ResponseInfo对象
    """
    ret,info=put_data(token,save_file_name,source_file.stream)

    print(type(info.status_code),info)
    if info.status_code==200:
        return domain_prefix+save_file_name
    return None


# -*- coding: utf-8 -*-

from __future__ import print_function, division

import argparse

# import neptune
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
from torchvision import datasets, transforms
import torch.backends.cudnn as cudnn
import matplotlib
import random
import scipy.io

matplotlib.use('agg')
import matplotlib.pyplot as plt
# from PIL import Image
import copy
import time
import os
from model import *
from random_erasing import RandomErasing
from autoaugment import ImageNetPolicy, CIFAR10Policy
import yaml
import math
from shutil import copyfile
from utils import *
import numpy as np
from image_folder import SatData, DroneData, ImageFolder_selectID, ImageFolder_expandID, customData
import wandb

# print(torch.version.cuda,"torch.version.cuda")
# 11.1 torch.version.cuda

version = torch.__version__
# fp16
try:
    from apex.fp16_utils import *
    from apex import amp, optimizers
except ImportError:  # will be 3.x series
    print(
        'This is not an error. If you want to use low precision, i.e., fp16, please install the apex with cuda support (https://github.com/NVIDIA/apex) and update pytorch to 1.0')
######################################################################
# Options
# --------
parser = argparse.ArgumentParser(description='Training')
parser.add_argument('--gpu_ids', default='0', type=str, help='gpu_ids: e.g. 0  0,1,2  0,2')
parser.add_argument('--name', default='debug', type=str, help='output model name')
parser.add_argument('--pool', default='avg', type=str, help='pool avg')
parser.add_argument('--data_dir', default='/home/lihaoran/BJDD_datesets/datesets/University-Release/train', type=str,
                    help='training dir path')
parser.add_argument('--train_all', action='store_true', help='use all training data')
parser.add_argument('--color_jitter', action='store_true', help='use color jitter in training')
parser.add_argument('--batchsize', default=8, type=int, help='batchsize')
parser.add_argument('--stride', default=1, type=int, help='stride')
parser.add_argument('--pad', default=10, type=int, help='padding')
parser.add_argument('--h', default=256, type=int, help='height')
parser.add_argument('--w', default=256, type=int, help='width')
parser.add_argument('--views', default=2, type=int, help='the number of views')
parser.add_argument('--erasing_p', default=0, type=float, help='Random Erasing probability, in [0,1]')
parser.add_argument('--use_dense', action='store_true', help='use densenet121')
parser.add_argument('--use_NAS', action='store_true', help='use NAS')
parser.add_argument('--warm_epoch', default=0, type=int, help='the first K epoch that needs warm up')
parser.add_argument('--lr', default=0.01, type=float, help='learning rate')
parser.add_argument('--moving_avg', default=1.0, type=float, help='moving average')
parser.add_argument('--droprate', default=0.75, type=float, help='drop rate')
parser.add_argument('--DA', action='store_true', help='use Color Data Augmentation')
parser.add_argument('--resume', action='store_true', help='use resume trainning')
parser.add_argument('--share', action='store_true', help='share weight between different view')
parser.add_argument('--extra_Google', action='store_true', help='using extra noise Google')
parser.add_argument('--LPN', action='store_true', help='use LPN')
parser.add_argument('--decouple', action='store_true', help='use decouple')
parser.add_argument('--block', default=4, type=int, help='the num of block')
parser.add_argument('--scale', default=1 / 32, type=float, metavar='S', help='scale the loss')
parser.add_argument('--lambd', default=3.9e-3, type=float, metavar='L', help='weight on off-diagonal terms')
parser.add_argument('--g', default=0.9, type=float, metavar='L', help='weight on loss and deloss')
parser.add_argument('--t', default=4.0, type=float, metavar='L', help='temperature of conv matrix')
parser.add_argument('--experiment_name', default='debug', type=str, help='log dir name')
parser.add_argument('--adam', action='store_true', help='using adam optimization')
parser.add_argument('--seed', default=0, type=int, help='random seed')
parser.add_argument('--balance', action='store_true', help='using balance sampler')
parser.add_argument('--select_id', action='store_true', help='select id')
parser.add_argument('--multi_image', action='store_true', help='only inputs3 + inputs3_s training')
parser.add_argument('--expand_id', action='store_true', help='expand id')
parser.add_argument('--dro_lead', action='store_true', help='drone leading sampling')
parser.add_argument('--sat_lead', action='store_true', help='satellite leading sampling')
parser.add_argument('--normal', action='store_true', help='normal training')
parser.add_argument('--only_decouple', action='store_true', help='do not use balance losss')
parser.add_argument('--e1', default=1, type=int, help='the exponent of on diag')
parser.add_argument('--e2', default=1, type=int, help='the exponent of off diag')
parser.add_argument('--swin', action='store_true', help='using swin as backbone')
parser.add_argument('--fp16', action='store_true',
                    help='use float16 instead of float32, which will save about 50% memory')
parser.add_argument('--test_dir', default='/home/lihaoran/BJDD_datesets/datesets/University-Release/test',
                    help='use float16 instead of float32, which will save about 50% memory')
parser.add_argument('--modelNum', type=int, help='用来区分模型')
parser.add_argument('--val_batchsize', default=128, type=int, help='batchsize')
parser.add_argument('--SAM', type=int, default=0, help='用来区分模型')
parser.add_argument('--infonce', type=int, default=1, help='采用infonce损失来平衡')
parser.add_argument('--Twozerothree', action='store_true', help='采用infonce损失来平衡')
opt = parser.parse_args()

os.environ["WANDB_API_KEY"] = 'c67d9a2bf298e65e8717d5c693270e77d117bddb'
os.environ["WANDB_MODE"] = "offline"
wandb.init(project="DWDR", name=opt.name)

def seed_torch(seed=opt.seed):
    # random.seed(seed)
    seed = 1234
    # print("1111111")
    # os.environ['PYTHONHASHSEED'] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    # np.random.seed(seed)
    # torch.manual_seed(seed)
    # torch.cuda.manual_seed(seed)
    # torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    # random.seed(seed)  # Python random module.
    # torch.backends.cudnn.benchmark = True
    # torch.backends.cudnn.deterministic = True
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)
    # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
    # torch.use_deterministic_algorithms(False)

print('random seed---------------------:', opt.seed)
seed_torch(opt.seed)

if opt.resume:
    model, opt, start_epoch = load_network(opt.name, opt)
else:
    start_epoch = 0

# debug
# opt.LPN=True
# opt.decouple = True

fp16 = opt.fp16
data_dir = opt.data_dir
test_dir = opt.test_dir
name = opt.name
str_ids = opt.gpu_ids.split(',')
gpu_ids = []
for str_id in str_ids:
    gid = int(str_id)
    if gid >= 0:
        gpu_ids.append(gid)

# set gpu ids
if len(gpu_ids) > 1:
    os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids
    cudnn.enabled = True
    cudnn.benchmark = True
    print('-----------------------------use multi-GPU',gpu_ids)

else:
    os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, gpu_ids))
    cudnn.benchmark = True
    testGpu = ''.join(map(str, gpu_ids))
print('---------------Pool Strategy------------:', opt.pool)
######################################################################
# Load Data
# ---------
#

transform_train_list = [
    # transforms.RandomResizedCrop(size=(opt.h, opt.w), scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC)
    transforms.Resize((opt.h, opt.w), interpolation=3),
    transforms.Pad(opt.pad, padding_mode='edge'),
    transforms.RandomCrop((opt.h, opt.w)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]

transform_satellite_list = [
    transforms.Resize((opt.h, opt.w), interpolation=3),
    transforms.Pad(opt.pad, padding_mode='edge'),
    transforms.RandomAffine(90),
    transforms.RandomCrop((opt.h, opt.w)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]

transform_val_list = [
    transforms.Resize(size=(opt.h, opt.w), interpolation=3),  # Image.BICUBIC
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]

if opt.erasing_p > 0:
    transform_train_list = transform_train_list + [RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0])]

if opt.color_jitter:
    transform_train_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1,
                                                   hue=0)] + transform_train_list
    transform_satellite_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1,
                                                       hue=0)] + transform_satellite_list

if opt.DA:
    transform_train_list = [ImageNetPolicy()] + transform_train_list

print(transform_train_list)
data_transforms = {
    'train': transforms.Compose(transform_train_list),
    'val': transforms.Compose(transform_val_list),
    'satellite': transforms.Compose(transform_satellite_list)
}

transform_move_list = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

train_all = ''
if opt.train_all:
    train_all = '_all'

image_datasets = {}
if opt.expand_id:
    print('--------------------expand id-----------------------')
    image_datasets['satellite'] = ImageFolder_expandID(os.path.join(data_dir, 'satellite'),
                                                       transform=data_transforms['satellite'])
else:
    image_datasets['satellite'] = SatData(data_dir, data_transforms['satellite'], data_transforms['train'])

if opt.select_id:
    print('--------------------select id-----------------------')
    image_datasets['drone'] = ImageFolder_selectID(os.path.join(data_dir, 'drone'), transform=data_transforms['train'])
else:
    image_datasets['drone'] = DroneData(data_dir, data_transforms['train'], data_transforms['satellite'])

def _init_fn(worker_id):
    np.random.seed(int(opt.seed) + worker_id)

dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=opt.batchsize,
                                              shuffle=True, num_workers=8, pin_memory=False, worker_init_fn=_init_fn)
               # 8 workers may work faster
               for x in ['satellite', 'drone']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['satellite', 'drone']}
class_names = image_datasets['satellite'].classes
print(dataset_sizes)

test_imgDatasets = {x: datasets.ImageFolder(os.path.join(test_dir, x), data_transforms['val']) for x in
                    ['gallery_satellite', 'gallery_drone']}
for x in ['query_satellite', 'query_drone']:
    test_imgDatasets[x] = customData(os.path.join(test_dir, x), data_transforms['val'], rotate=0)
Val_dataloaders = {x: torch.utils.data.DataLoader(test_imgDatasets[x], batch_size=opt.val_batchsize,
                                                  shuffle=False, num_workers=8) for x in
                   ['gallery_satellite', 'gallery_drone', 'query_satellite', 'query_drone']}

use_gpu = torch.cuda.is_available()

######################################################################
# Training the model
# ------------------
#
# Now, let's write a general function to train a model. Here, we will
# illustrate:
#
# -  Scheduling the learning rate
# -  Saving the best model
#
# In the following, parameter ``scheduler`` is an LR scheduler object from
# ``torch.optim.lr_scheduler``.

y_loss = {}  # loss history
y_loss['train'] = []
y_loss['val'] = []
y_err = {}
y_err['train'] = []
y_err['val'] = []

def compute_mAP(index, good_index, junk_index):
    ap = 0
    cmc = torch.IntTensor(len(index)).zero_()
    if good_index.size == 0:  # if empty
        cmc[0] = -1
        return ap, cmc

    # remove junk_index
    mask = np.in1d(index, junk_index, invert=True)
    index = index[mask]

    # find good_index index
    ngood = len(good_index)
    mask = np.in1d(index, good_index)
    rows_good = np.argwhere(mask == True)
    rows_good = rows_good.flatten()

    cmc[rows_good[0]:] = 1
    for i in range(ngood):
        d_recall = 1.0 / ngood
        precision = (i + 1) * 1.0 / (rows_good[i] + 1)
        if rows_good[i] != 0:
            old_precision = i * 1.0 / rows_good[i]
        else:
            old_precision = 1.0
        ap = ap + d_recall * (old_precision + precision) / 2

    return ap, cmc

def evaluate(qf, ql, gf, gl):
    query = qf.view(-1, 1)
    # print(query.shape)
    score = torch.mm(gf, query)
    score = score.squeeze(1).cpu()
    score = score.numpy()
    # predict index
    index = np.argsort(score)  # from small to large
    index = index[::-1]
    # index = index[0:2000]
    # good index
    query_index = np.argwhere(gl == ql)
    good_index = query_index
    # print(good_index)
    # print(index[0:10])
    junk_index = np.argwhere(gl == -1)

    CMC_tmp = compute_mAP(index, good_index, junk_index)
    return CMC_tmp

def which_view(name):
    if 'satellite' in name:
        return 1
    elif 'street' in name:
        return 2
    elif 'drone' in name:
        return 3
    else:
        print('unknown view')
    return -1

def fliplr(img):
    '''flip horizontal'''
    inv_idx = torch.arange(img.size(3) - 1, -1, -1).long()  # N x C x H x W
    img_flip = img.index_select(3, inv_idx)
    return img_flip

def get_id(img_path):
    camera_id = []
    labels = []
    paths = []
    for path, v in img_path:
        # print(path, v)
        folder_name = os.path.basename(os.path.dirname(path))
        labels.append(int(folder_name))
        paths.append(path)
    return labels, paths

def extract_feature(model, dataloaders, view_index=1):
    features = torch.FloatTensor()
    count = 0
    for data in dataloaders:
        img, label = data
        n, c, h, w = img.size()
        count += n
        # print(count)
        ff = torch.FloatTensor(n, 512).zero_().cuda()
        # if opt.swin:
        #     ff = torch.FloatTensor(n,1024).zero_().cuda()
        if opt.LPN:
            # ff = torch.FloatTensor(n,2048,6).zero_().cuda()
            ff = torch.FloatTensor(n, 512, opt.block).zero_().cuda()

        for i in range(2):
            if (i == 1):
                img = fliplr(img)
            input_img = Variable(img.cuda())
            outputs = model(input_img)
            # print(outputs.shape)
            if opt.decouple:
                ff += outputs[0]
            elif opt.infonce == 1:
                ff += outputs[0]
            else:
                ff += outputs
        # norm feature
        if opt.LPN:
            # feature size (n,2048,6)
            # 1. To treat every part equally, I calculate the norm for every 2048-dim part feature.
            # 2. To keep the cosine score==1, sqrt(6) is added to norm the whole feature (2048*6).
            fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) * np.sqrt(opt.block)
            ff = ff.div(fnorm.expand_as(ff))
            ff = ff.view(ff.size(0), -1)
        else:
            fnorm = torch.norm(ff, p=2, dim=1, keepdim=True)
            ff = ff.div(fnorm.expand_as(ff))

        features = torch.cat((features, ff.data.cpu()), 0)
    return features

# work channel loss
def off_diagonal(x):
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()

def decouple_loss(y1, y2, scale_loss, lambd):
    batch_size = y1.size(0)
    c = y1.T @ y2
    c.div_(batch_size)
    on_diag = torch.diagonal(c)
    p_on = (1 - on_diag) / 2
    on_diag = torch.pow(p_on, opt.e1) * torch.pow(torch.diagonal(c).add_(-1), 2)
    on_diag = on_diag.sum().mul(scale_loss)

    off_diag = off_diagonal(c)
    p_off = torch.abs(off_diag)
    off_diag = torch.pow(p_off, opt.e2) * torch.pow(off_diagonal(c), 2)
    off_diag = off_diag.sum().mul(scale_loss)
    loss = on_diag + off_diag * lambd
    return loss, on_diag, off_diag * lambd

def one_LPN_output(outputs, labels, criterion, block):
    # part = {}
    sm = nn.Softmax(dim=1)
    num_part = block
    score = 0
    loss = 0
    for i in range(num_part):
        part = outputs[i]
        score += sm(part)
        loss += criterion(part, labels)

    _, preds = torch.max(score.data, 1)

    return preds, loss

def one_info_output(Soutputs, Doutputs, labels, criterion, block):
    num_part = block
    for i in range(num_part):
        Dpart = Soutputs[i]  # 2,701
        Spart = Doutputs[i]  # 2,701
        s_norm = F.normalize(Dpart, dim=1)
        d_norm = F.normalize(Spart, dim=1)
        features = torch.cat([s_norm, d_norm], dim=1)  # 2,701

def val(model, epoch, ):
    config_path = os.path.join('./model', opt.name, 'opts.yaml')
    with open(config_path, 'r') as stream:
        config = yaml.safe_load(stream)

    model_test = copy.deepcopy(model)

    if opt.LPN:
        if len(gpu_ids) > 1:
            for i in range(opt.block):
                cls_name = 'classifier' + str(i)
                c = getattr(model_test.module, cls_name)
                c.classifier = nn.Sequential()
        # model = three_view_net_test(model)

        else:
            for i in range(opt.block):
                cls_name = 'classifier' + str(i)
                c = getattr(model_test, cls_name)
                c.classifier = nn.Sequential()
    else:
        model_test.classifier.classifier = nn.Sequential()
        # model.classifier = nn.Sequential()
    model_test = model_test.eval()
    model_test = model_test.cuda()
    query_name = 'query_drone'
    gallery_name = 'gallery_satellite'
    which_gallery = which_view(gallery_name)
    which_query = which_view(query_name)
    print('%d -> %d:' % (which_query, which_gallery))
    gallery_path = test_imgDatasets[gallery_name].imgs
    f = open('gallery_name.txt', 'w')
    for p in gallery_path:
        f.write(p[0] + '\n')
    query_path = test_imgDatasets[query_name].imgs
    f = open('query_name.txt', 'w')
    for p in query_path:
        f.write(p[0] + '\n')
    gallery_label, gallery_path = get_id(gallery_path)
    query_label, query_path = get_id(query_path)

    with torch.no_grad():
        query_feature = extract_feature(model_test, Val_dataloaders[query_name], which_query)
        gallery_feature = extract_feature(model_test, Val_dataloaders[gallery_name], which_gallery)

        time_elapsed = time.time() - since
        print('Test complete in {:.0f}m {:.0f}s'.format(
            time_elapsed // 60, time_elapsed % 60))

        # Save to Matlab for check
        result = {'gallery_f': gallery_feature.numpy(), 'gallery_label': gallery_label, 'gallery_path': gallery_path,
                  'query_f': query_feature.numpy(), 'query_label': query_label, 'query_path': query_path}
        scipy.io.savemat('pytorch_result.mat', result)

        print(opt.name)
        result = './model/%s/result.txt' % opt.name
        # os.system(
        #     'CUDA_VISIBLE_DEVICES=%d python evaluate_gpu.py --name %s | tee -a %s' % (testGpu, opt.name, result))

        result = scipy.io.loadmat('pytorch_result.mat')
        query_feature = torch.FloatTensor(result['query_f'])
        query_label = result['query_label'][0]
        gallery_feature = torch.FloatTensor(result['gallery_f'])
        gallery_label = result['gallery_label'][0]
        multi = os.path.isfile('multi_query.mat')

        if multi:
            m_result = scipy.io.loadmat('multi_query.mat')
            mquery_feature = torch.FloatTensor(m_result['mquery_f'])
            mquery_label = m_result['mquery_label'][0]
            mquery_feature = mquery_feature.cuda()

        query_feature = query_feature.cuda()
        gallery_feature = gallery_feature.cuda()

        CMC = torch.IntTensor(len(gallery_label)).zero_()
        ap = 0.0

        for i in range(len(query_label)):
            ap_tmp, CMC_tmp = evaluate(query_feature[i], query_label[i], gallery_feature, gallery_label)
            if CMC_tmp[0] == -1:
                continue
            CMC = CMC + CMC_tmp
            ap += ap_tmp

        CMC = CMC.float()
        CMC = CMC / len(query_label)  # average CMC
        print(round(len(gallery_label) * 0.01))
        acc1 = CMC[0] * 100
        ap1 = ap / len(query_label) * 100
        print('Recall@1:%.2f Recall@5:%.2f Recall@10:%.2f Recall@top1:%.2f AP:%.2f' % (
            acc1, CMC[4] * 100, CMC[9] * 100, CMC[round(len(gallery_label) * 0.01)] * 100, ap1))

        # 放到wandb中去

        # multiple-query
        CMC = torch.IntTensor(len(gallery_label)).zero_()
        ap = 0.0
        if multi:
            for i in range(len(query_label)):
                mquery_index1 = np.argwhere(mquery_label == query_label[i])
                mquery_index2 = np.argwhere(mquery_cam == query_cam[i])
                mquery_index = np.intersect1d(mquery_index1, mquery_index2)
                mq = torch.mean(mquery_feature[mquery_index, :], dim=0)
                ap_tmp, CMC_tmp = evaluate(mq, query_label[i], query_cam[i], gallery_feature, gallery_label,
                                           gallery_cam)
                if CMC_tmp[0] == -1:
                    continue
                CMC = CMC + CMC_tmp
                ap += ap_tmp
                # print(i, CMC_tmp[0])
            CMC = CMC.float()
            CMC = CMC / len(query_label)  # average CMC
            print('multi Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f' % (CMC[0], CMC[4], CMC[9], ap / len(query_label)))

        return acc1, ap1

def train_model(model, model_test, criterion, optimizer, scheduler, epoch, warm_up, warm_iteration, num_epochs=25):
    epoch = epoch + start_epoch
    print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    print('-' * 10)

    # Each epoch has a training and validation phase
    for phase in ['train']:
        if phase == 'train':
            model.train(True)  # Set model to training mode
        else:
            model.train(False)  # Set model to evaluate mode

        running_loss = 0.0
        running_corrects = 0.0
        running_corrects3 = 0.0
        ins_loss = 0.0
        dec_loss = 0.0
        on_loss = 0.0
        off_loss = 0.0
        lossinfo1 = 0.0
        lossinfo2 = 0.0
        # Iterate over data.
        for data, data3 in zip(dataloaders['satellite'], dataloaders['drone']):
            # get the inputs
            inputs, inputs_d, labels = data
            # print(inputs.shape,"inputs")  # torch.Size([8, 3, 224, 224])
            # print(inputs_d.shape,"inputs_d")  # torch.Size([8, 3, 224, 224])
            # print(labels,"labels")  # tensor([360, 142, 532, 398, 322, 141, 199, 609])
            inputs3, inputs3_s, labels3 = data3
            now_batch_size, c, h, w = inputs.shape
            if now_batch_size < opt.batchsize:  # skip the last batch
                continue
            if use_gpu:
                if opt.normal:
                    inputs = Variable(inputs.cuda().detach())
                    inputs3 = Variable(inputs3.cuda().detach())
                    labels = Variable(labels.cuda().detach())
                    labels3 = Variable(labels3.cuda().detach())
                else:
                    inputs = Variable(inputs.cuda().detach())
                    inputs_d = Variable(inputs_d.cuda().detach())
                    inputs3 = Variable(inputs3.cuda().detach())
                    inputs3_s = Variable(inputs3_s.cuda().detach())
                    labels = Variable(labels.cuda().detach())
                    labels3 = Variable(labels3.cuda().detach())

            else:
                inputs, labels = Variable(inputs), Variable(labels)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward
            if opt.decouple:
                if opt.infonce == 1:
                    outs_c, outs_f, outs_info = model(inputs)
                else:
                    outs_c, outs_f = model(inputs)
            else:
                if opt.infonce == 1:
                    outs_c, outs_info = model(inputs)
                else:
                    outs_c = model(inputs)
            if opt.balance:
                if opt.decouple:
                    if opt.infonce == 1:
                        outs_d_c, outs_d_f, outs_d_info = model(inputs_d)
                    else:
                        outs_d_c, outs_d_f = model(inputs_d)
                else:
                    if opt.infonce == 1:
                        outs_d_c, outs_d_info = model(inputs_d)
                    else:
                        outs_d_c = model(inputs_d)

            if opt.decouple:
                if opt.infonce == 1:
                    outd_c, outs3_f, outd_info = model(inputs3)
                else:
                    outd_c, outs3_f = model(inputs3)
            else:
                if opt.infonce == 1:
                    outd_c, outd_info = model(inputs3)
                else:
                    outd_c = model(inputs3)
            if opt.balance:
                if opt.decouple:
                    if opt.infonce == 1:
                        outs3_s_c, outs3_s_f, outs3_s_info = model(inputs3_s)
                    else:
                        outs3_s_c, outs3_s_f = model(inputs3_s)
                else:
                    if opt.infonce == 1:
                        outs3_s_c, outs3_s_info = model(inputs3_s)
                    else:
                        outs3_s_c = model(inputs3_s)
            # calculate loss
            if opt.LPN:
                if opt.balance:
                    # print('--------------------- using data balance---------------------------')
                    if opt.only_decouple:
                        print('--------------------- only decouple---------------------------')
                        preds, loss = one_LPN_output(outs_c, labels, criterion, opt.block)
                        preds3, loss3 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                        loss = loss + loss3
                    else:

                        preds, loss = one_LPN_output(outs_c, labels, criterion, opt.block)
                        _, loss_d = one_LPN_output(outs_d_c, labels, criterion, opt.block)
                        loss = loss + loss_d
                        preds3, loss3 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                        _, loss3_s = one_LPN_output(outs3_s_c, labels3, criterion, opt.block)
                        loss3 = loss3 + loss3_s
                        loss = (loss + loss3) / 2

                else:

                    preds, loss = one_LPN_output(outs_c, labels, criterion, opt.block)
                    preds3, loss3 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                    loss = loss + loss3

                if opt.decouple:

                    if opt.balance:
                        deloss1, on, off = decouple_loss(outs_f, outs_d_f, opt.scale, opt.lambd)
                        deloss2, on1, off1 = decouple_loss(outs3_s_f, outs3_f, opt.scale, opt.lambd)
                        deloss = (deloss1 + deloss2) / 2
                        # deloss = deloss2
                        on = (on + on1) / 2
                        off = (off + off1) / 2
                        insloss = loss
                        loss = opt.g * insloss + (1 - opt.g) * deloss

                # 在这里需要加入infonce的loss计算方法  在使用balance的时候需要进行双向的加权（一个正常的loss一个color变化的loss）

                if opt.infonce == 1 and opt.balance :  # 默认使用
                    # 正常图片下的infonce
                    sate = F.normalize(outs_info, dim=1)
                    drone = F.normalize(outd_info, dim=1)
                    sate_ = F.normalize(outs_d_info, dim=1)
                    drone_ = F.normalize(outs3_s_info, dim=1)

                    features1 = torch.cat([sate.unsqueeze(1), sate_.unsqueeze(1)], dim=1)
                    features2 = torch.cat([drone.unsqueeze(1), drone_.unsqueeze(1)], dim=1)

                    loss_info = infonce(features1, labels)
                    loss = loss + loss_info

                    loss_info1 = infonce(features2, labels3)
                    loss = loss + loss_info1
            else:
                pass

            if epoch < opt.warm_epoch and phase == 'train':
                warm_up = min(1.0, warm_up + 0.9 / warm_iteration)
                loss *= warm_up

            if phase == 'train':
                if fp16:  # we use optimier to backward loss
                    loss.backward()
                else:
                    loss.backward()
                if opt.SAM != 1:
                    optimizer.step()
                else:
                    optimizer.first_step(zero_grad=True)
                    # sam中计算第二次梯度
                    if opt.balance:

                        if opt.decouple:  # 同时使用balance和decouple的情况下

                            if opt.infonce == 1:
                                outs_c, outs_f, outs_info = model(inputs)
                                outd_c, outs3_f, outd_info = model(inputs3)
                                outs_d_c, outs_d_f, outs_d_info = model(inputs_d)
                                outs3_s_c, outs3_s_f, outs3_s_info = model(inputs3_s)

                            else:
                                outs_c, outs_f = model(inputs)
                                outd_c, outs3_f = model(inputs3)
                                outs3_s_c, outs3_s_f = model(inputs3_s)
                                outs_d_c, outs_d_f = model(inputs_d)

                            preds, loss = one_LPN_output(outs_c, labels, criterion, opt.block)
                            _, loss_d = one_LPN_output(outs_d_c, labels, criterion, opt.block)
                            loss = loss + loss_d
                            preds3, loss3 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                            _, loss3_s = one_LPN_output(outs3_s_c, labels3, criterion, opt.block)
                            loss3 = loss3 + loss3_s
                            loss = (loss + loss3) / 2
                            deloss1, on, off = decouple_loss(outs_f, outs_d_f, opt.scale, opt.lambd)
                            deloss2, on1, off1 = decouple_loss(outs3_s_f, outs3_f, opt.scale, opt.lambd)
                            deloss = (deloss1 + deloss2) / 2
                            # deloss = deloss2
                            on = (on + on1) / 2
                            off = (off + off1) / 2
                            insloss = loss
                            loss = opt.g * insloss + (1 - opt.g) * deloss

                        else:

                            if opt.infonce == 1:

                                outd_c, outd_info = model(inputs3)
                                outs_c, outs_info = model(inputs)
                                outs_d_c, outs_d_info = model(inputs_d)
                                outs3_s_c, outs3_s_info = model(inputs3_s)

                            else:

                                outs_c = model(inputs)
                                outd_c = model(inputs3)
                                outs3_s_c = model(inputs3_s)
                                outs_d_c = model(inputs_d)

                            preds, loss = one_LPN_output(outs_c, labels, criterion, opt.block)
                            _, loss_d = one_LPN_output(outs_d_c, labels, criterion, opt.block)
                            loss = loss + loss_d
                            preds3, loss3 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                            _, loss3_s = one_LPN_output(outs3_s_c, labels3, criterion, opt.block)
                            loss3 = loss3 + loss3_s
                            loss = (loss + loss3) / 2

                        # 使用infonce的情况下计算新一轮的损失？

                    else:

                        # 在不适用balance和decouple的时候去使用infonce

                        if opt.infonce == 1:
                            outs_c, outs_info = model(inputs)
                            outd_c, outd_info = model(inputs3)
                        else:
                            outs_c = model(inputs)
                            outd_c = model(inputs3)

                        preds, lossmin3 = one_LPN_output(outs_c, labels, criterion, opt.block)
                        preds3, lossmin4 = one_LPN_output(outd_c, labels3, criterion, opt.block)
                        loss = (lossmin3 + lossmin4)

                    if opt.infonce == 1 and opt.balance:
                        # 最后统一计算损失
                        sate = F.normalize(outs_info, dim=1)
                        drone = F.normalize(outd_info, dim=1)
                        sate_ = F.normalize(outs_d_info, dim=1)
                        drone_ = F.normalize(outs3_s_info, dim=1)

                        features1 = torch.cat([sate.unsqueeze(1), sate_.unsqueeze(1)], dim=1)
                        features2 = torch.cat([drone.unsqueeze(1), drone_.unsqueeze(1)], dim=1)

                        loss_info = infonce(features1, labels)
                        loss = loss + loss_info

                        loss_info1 = infonce(features2, labels3)
                        loss = loss + loss_info1

                    if epoch < opt.warm_epoch and phase == 'train':
                        loss *= warm_up
                    # print(loss2, "2")
                    if fp16:  # we use optimier to backward loss
                        loss.backward()
                    else:
                        loss.backward()
                    optimizer.second_step(zero_grad=True)
                ##########
            if opt.moving_avg < 1.0:
                update_average(model_test, model, opt.moving_avg)

            # statistics

            running_loss += loss.item() * now_batch_size
            if opt.decouple:
                ins_loss += insloss.item() * now_batch_size
                dec_loss += deloss.item() * now_batch_size
                on_loss += on.item() * now_batch_size
                off_loss += off.item() * now_batch_size

            running_corrects += float(torch.sum(preds == labels.data))
            running_corrects3 += float(torch.sum(preds3 == labels3.data))

        epoch_loss = running_loss / dataset_sizes['satellite']
        epoch_acc = running_corrects / dataset_sizes['satellite']
        epoch_acc3 = running_corrects3 / dataset_sizes['satellite']

        if opt.decouple:
            epoch_ins_loss = ins_loss / dataset_sizes['satellite']
            epoch_dec_loss = dec_loss / dataset_sizes['satellite']
            epoch_on_loss = on_loss / dataset_sizes['satellite']
            epoch_off_loss = off_loss / dataset_sizes['satellite']

        if opt.infonce == 1:
            lossinfo1 += loss_info.item() * now_batch_size
            # lossinfo2 += loss_info1.item() * now_batch_size
            epoch_loss_info1 = lossinfo1 / dataset_sizes['satellite']
            # epoch_loss_info2 = lossinfo2 / dataset_sizes['satellite']

        if opt.decouple:
            print(
                '{} Loss: {:.4f} Satellite_Acc: {:.4f} Drone_Acc: {:.4f}, On_Loss: {:.4f}, Off_Loss: {:.4f},'.format(
                    phase, epoch_loss, epoch_acc, epoch_acc3, epoch_on_loss, epoch_off_loss))

        if opt.infonce == 1:
            print('{} Loss: {:.4f} Satellite_Acc: {:.4f} Drone_Acc: {:.4f} infoloss1: {:.4f} infoloss2: {:.4f}'.format(
                phase, epoch_loss, epoch_acc,
                epoch_acc3, epoch_loss_info1, 0.00))

        else:
            print('{} Loss: {:.4f} Satellite_Acc: {:.4f} Drone_Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc,
                                                                                epoch_acc3))

        y_loss[phase].append(epoch_loss)
        y_err[phase].append(1.0 - epoch_acc)

        # saving last model:
        if phase == 'train':
            scheduler.step()
        if epoch + 1 == num_epochs and len(gpu_ids) > 1:
            save_network(model.module, opt.name, epoch)
        elif epoch + 1 > 100 and (epoch + 1) % 10 == 0:
            save_network(model, opt.name, epoch)
        # draw_curve(epoch)

    if epoch % 4 == 1:
        wandb.log({
            'Step': epoch + 1,
            'Loss': epoch_loss,
            'Satellite_Acc': epoch_acc,
            'Drone_Acc': epoch_acc3,
        })

    return model, warm_up

######################################################################
# Draw Curve
# ---------------------------
x_epoch = []
fig = plt.figure()
ax0 = fig.add_subplot(121, title="loss")
ax1 = fig.add_subplot(122, title="top1err")

def draw_curve(current_epoch):
    x_epoch.append(current_epoch)
    ax0.plot(x_epoch, y_loss['train'], 'bo-', label='train')
    ax0.plot(x_epoch, y_loss['val'], 'ro-', label='val')
    ax1.plot(x_epoch, y_err['train'], 'bo-', label='train')
    ax1.plot(x_epoch, y_err['val'], 'ro-', label='val')
    if current_epoch == 0:
        ax0.legend()
        ax1.legend()
    fig.savefig(os.path.join('./model', name, 'train.jpg'))

######################################################################
# Finetuning the convnet
# ----------------------

# Load a pretrainied model and reset final fully connected layer.
#
if opt.LPN:
    # model = ft_net_LPN(len(class_names), droprate=opt.droprate, stride=opt.stride, pool=opt.pool, block=opt.block,
    #                    decouple=opt.decouple)
    model = CSWinTransv2_threeIn(len(class_names), droprate=opt.droprate, decouple=opt.decouple, infonce=opt.infonce)

    if opt.Twozerothree:
        from model_203 import *

        model = CSWinTrans_attention(len(class_names), droprate=opt.droprate, decouple=opt.decouple,
                                     infonce=opt.infonce)

    # model = CSWinTrans_twoStage(len(class_names), droprate=opt.droprate)
# elif opt.swin:
#     model = ft_net_swin(len(class_names), droprate=opt.droprate, decouple=opt.decouple)
# else:
#     model = mainModule(len(class_names), droprate=opt.droprate, decouple=opt.decouple)
#     # model = TransCnn_50(len(class_names), droprate=opt.droprate)
#     # model = SAIG_Shallow(img_size=224)

# model = two_view_net()

opt.nclasses = len(class_names)
print('nclass--------------------:', opt.nclasses)
# print(model)
# For resume:
if start_epoch >= 40:
    opt.lr = opt.lr * 0.1

if not opt.LPN:
    model = model.cuda()

    params = [{"params": model.get_1x_lr_params(), "lr": opt.lr * 0.1},
              {"params": model.get_10x_lr_params(), "lr": opt.lr},
              # {"params": model.model.get_add_lr_params(), "lr": opt.lr * 1.5 }
              ]

    # ignored_params = list(map(id, model.classifier.parameters()))
    # # ignored_params.append(map(id,model.model.t2tswin.parameters()))
    # base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
    optimizer_ft = optim.SGD(
        params,
        #     [
        #     {'params': base_params, 'lr': 0.1 * opt.lr},
        #     {'params': model.classifier.parameters(), 'lr': opt.lr},
        #     # {'params':model.model.t2tswin.parameters(),'lr':opt.lr}
        # ],
        weight_decay=5e-4, momentum=0.9, nesterov=True)

    # print(ignored_params)

    # print(base_params,"base_params")
    # print(model.model.classifier,"model.model.classifier.parameters()")

else:
    # ignored_params = list(map(id, model.model.fc.parameters() ))
    if len(gpu_ids) > 1:
        print(gpu_ids,"gpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_idsgpu_ids")
        model = torch.nn.DataParallel(model,device_ids=[0,1]).cuda()

        ignored_params = list()
        for i in range(opt.block):
            cls_name = 'classifier' + str(i)
            c = getattr(model.module, cls_name)
            ignored_params += list(map(id, c.parameters()))

        base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())

        optim_params = [{'params': base_params, 'lr': 0.1 * opt.lr}]
        for i in range(opt.block):
            cls_name = 'classifier' + str(i)
            c = getattr(model.module, cls_name)
            optim_params.append({'params': c.parameters(), 'lr': opt.lr})

        # optim_params = [{"params": model.module.get_1x_lr_params(), "lr": opt.lr * 0.1},
        #                 {"params": model.module.get_10x_lr_params(), "lr": opt.lr},
        #                 # {"params": model.model.get_add_lr_params(), "lr": opt.lr * 1.5 }
        #                 ]

    else:

        model = model.cuda()

        print('---------------------use one gpu-----------------------')
        ignored_params = list()
        # ignored_params += list(map(id, model.rdim.parameters() ))
        for i in range(opt.block):
            cls_name = 'classifier' + str(i)
            c = getattr(model, cls_name)
            ignored_params += list(map(id, c.parameters()))

        base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())

        optim_params = [{'params': base_params, 'lr': 0.1 * opt.lr}]
        # optim_params.append({'params': model.rdim.parameters(), 'lr': opt.lr})
        for i in range(opt.block):
            cls_name = 'classifier' + str(i)
            c = getattr(model, cls_name)
            optim_params.append({'params': c.parameters(), 'lr': opt.lr})

        # 学习率调整需要更改

        # optim_params = [{"params": model.get_1x_lr_params(), "lr": opt.lr * 0.1},
        #                 {"params": model.get_10x_lr_params(), "lr": opt.lr},
        #                 # {"params": model.model.get_add_lr_params(), "lr": opt.lr * 1.5 }
        #             ]

    # optimizer_ft = optim.SGD(optim_params, weight_decay=5e-4, momentum=0.9, nesterov=True)
    # base_optimizer = optim.SGD(optim_params, weight_decay=5e-4, momentum=0.9, nesterov=True)
    # base_optimizer = torch.optim.AdamW
    # optimizer_ft = SAM(optim_params, base_optimizer, lr=opt.lr, betas=(0.9, 0.999), weight_decay=5e-4, amsgrad=False,
    #                    adaptive=True, rho=2.5)

    infonce = SupConLoss(temperature=0.1)

    if opt.adam:
        optimizer_ft = optim.Adam(optim_params, opt.lr, weight_decay=5e-4)

    if opt.SAM == 1:
        base_optimizer = torch.optim.SGD
        optimizer_ft = SAM(optim_params, base_optimizer, lr=opt.lr, weight_decay=5e-4, momentum=0.9, nesterov=True)
        # base_optimizer = torch.optim.Adam
        # optimizer_ft = SAM(optim_params, base_optimizer, lr=opt.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, )
    else:
        optimizer_ft = optim.SGD(optim_params, weight_decay=5e-4, momentum=0.9, nesterov=True)

# Decay LR by a factor of 0.1 every 40 epochs
# exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=80, gamma=0.1)
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft, milestones=[120, 180, 210], gamma=0.1)

# exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft, milestones=[60,120,160], gamma=0.1)
# exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(optimizer_ft, T_max=120, eta_min=0.001)
######################################################################
# Train and evaluate
# ^^^^^^^^^^^^^^^^^^
#
# It should take around 1-2 hours on GPU.
#
# neptune.init('wtyu/decouple')
# neptune.create_experiment('LPN+norm(batch*512*4)')

if fp16:
    model, optimizer_ft = amp.initialize(model, optimizer_ft, opt_level="O1")

# if len(gpu_ids)>1:
#     model = torch.nn.DataParallel(model, [3,2]).cuda()
# else:
#     model = model.cuda()

criterion = nn.CrossEntropyLoss()

if opt.moving_avg < 1.0:
    model_test = copy.deepcopy(model)
    # num_epochs = 140
    num_epochs = 210
else:
    model_test = None
    # num_epochs = 120
    num_epochs = 210

warm_up = 0.1  # We start from the 0.1*lrRate
bestAcc = 0
bestAp = 0
bestEp = 0
since = time.time()

warm_iteration = round(dataset_sizes['satellite'] / opt.batchsize) * opt.warm_epoch  # first 5 epoch

for epoch in range(num_epochs - start_epoch):

    # train
    model, warm_up = train_model(model, model_test, criterion, optimizer_ft, exp_lr_scheduler, epoch, warm_up,
                                 warm_iteration,
                                 num_epochs=num_epochs)
    warm_up = warm_up

    dir_name = os.path.join('./model', name)
    if not opt.resume:
        if not os.path.isdir(dir_name):
            os.mkdir(dir_name)
        # record every run
        copyfile('./run_mul_gpu_view.sh', dir_name + '/run_mul_gpu_view.sh')
        copyfile('./train_info.py', dir_name + '/train_info.py')
        copyfile('./model.py', dir_name + '/model.py')
        # save opts
        with open('%s/opts.yaml' % dir_name, 'w') as fp:
            yaml.dump(vars(opt), fp, default_flow_style=False)

    # val
    # if epoch % 1 == 0:
    if epoch > 100 and epoch % 10 == 0:
        save_network(model, opt.name, epoch)
        acc1, ap1 = val(model, epoch)
        wandb.log({
            'acc@1': acc1,
            'ap': ap1,
            'Step': epoch + 1
        })
        if acc1 > bestAcc:
            bestAcc = acc1
            bestEp = epoch
        if ap1 > bestAp:
            bestAp = ap1

        print('BestRecall@1:%.2f, BestAP:%.2f,bestepoch:%.0f' % (bestAcc, bestAp, bestEp))

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('BestRecall@1:%.2f, BestAP:%.2f,bestepoch:%.0f' % (bestAcc, bestAp, bestEp))
    print("name:", opt.name)
    # if epoch_loss < best_loss:
    #     best_loss = epoch_loss
    #     best_epoch = epoch
    #     last_model_wts = model.state_dict()

    # time_elapsed = time.time() - since
    # print('Training complete in {:.0f}m {:.0f}s'.format(
    #     time_elapsed // 60, time_elapsed % 60))
    # print('Best val Acc: {:4f}'.format(best_acc))
    # model.load_state_dict(last_model_wts)
    # if len(gpu_ids)>1:
    #     save_network(model.module, opt.name, 'last')
    #     print('best_epoch:', best_epoch)
    # else:
    #     save_network(model, opt.name, 'last')
    #     print('best_epoch:', best_epoch)

# model = train_model(model, model_test, criterion, optimizer_ft, exp_lr_scheduler,
#                     num_epochs=num_epochs)
# neptune.stop()
print('BestRecall@1:%.2f, BestAP:%.2f' % (bestAcc, bestAp))
save_network(model, "best", epoch)