import Pretrained_supporting_functions as sf
import supporting_models as sm
import tensorflow as tf
from tensorflow import keras
import wandb
from wandb.keras import WandbCallback
import numpy as np
import matplotlib.pyplot as plt
import time
import tracemalloc
import os

#/vol/research/NOBACKUP/CVSSP/scratch_4weeks/ad00878/DBs/
#DBs/

def main():

    @tf.function
    def train_step(imgs,labels):
        with tf.GradientTape() as tape:
            preds = model(imgs,training=True)
            loss = loss_func(labels,preds)
        grads = tape.gradient(loss,model.trainable_variables)
        optimizer.apply_gradients(zip(grads,model.trainable_variables))
        train_loss(loss)
        train_acc_metric(labels,preds)
        train_prec_metric(labels,preds)
        train_rec_metric(labels,preds)
        return

    wandb.init(project='k_diversity',entity='adamdowse')
    config= {
        'db_path' : "DBs/",
        'ds_name' : "cifar10",
        'train_percent' : 0.1,
        'test_percent' : 0.1,
        'group' : 't9_hyperparam_search',
        'model_name' : 'Simple_CNN',
        'learning_rate' : wandb.config.learning_rate,
        'optimizer' : wandb.config.optimizer,
        'momentum' : 0,
        'random_db' : 'True',
        'batch_size' : wandb.config.batch_size,
        'data_aug' : '0', #0 = no data aug, 1 = data aug, 2 = data aug + noise
        'max_its' : 30000,
        'early_stop' : 5000,
        'mod_type' : 'Random_Bucket_full',
        'subset_type' : 'All', #Random_Bucket, Hard_Mining, All
        'train_type' : wandb.config.run_type, #SubMod, Random
        'activations_delay' : 4, #cannot be 0
        'k_percent' : 1, #percent of data to use for RB and HM
        'activation_layer_name' : 'penultimate_layer',
    }

    #Setup
    test_ds,ds_info,conn_path, train_ds = sf.setup_db(config)

    num_classes = ds_info.features['label'].num_classes
    class_names = ds_info.features['label'].names
    img_size = ds_info.features['image'].shape
    num_train_imgs = ds_info.splits['train[:'+str(int(config['train_percent']*100))+'%]'].num_examples

    #Model
    tf.keras.backend.clear_session()
    model = sm.select_model(config['model_name'],num_classes,img_size)
    model.build(img_size+(1,))
    model.summary()

    #Load pretrained weights TODO: make this a function

    #Loss
    loss_func = tf.keras.losses.CategoricalCrossentropy(from_logits=False)

    #Metrics
    train_loss = tf.keras.metrics.Mean(name='train_loss')
    train_acc_metric = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')
    train_prec_metric = tf.keras.metrics.Precision(name='train_precision')
    train_rec_metric = tf.keras.metrics.Recall(name='train_recall')

    #Optimizer
    if config['optimizer'] == 'Adam':
        optimizer = tf.keras.optimizers.Adam(learning_rate=config['learning_rate'], beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=False)
    elif config['optimizer'] == 'SGD':
        optimizer = tf.keras.optimizers.SGD(learning_rate=config['learning_rate'])
    elif config['optimizer'] == 'Momentum':
        optimizer = tf.keras.optimizers.SGD(learning_rate=config['learning_rate'],momentum=config['momentum'])
    else:
        print('Optimizer not recognised')

    #Data Generator
    train_DG = sf.SubModDataGen(conn_path,config)
    test_DG = sf.TestDataGen(test_ds, 50, num_classes)

    #Compile Model
    model.compile(optimizer=optimizer,loss=loss_func,metrics=['accuracy',tf.keras.metrics.Precision(),tf.keras.metrics.Recall()])

    #Wandb
    #sf.wandb_setup(config,disabled)
    logging_callback = WandbCallback(log_freq=1,save_model=False,save_code=False)

    #Training
    batch_num = 0
    early_stop_max = 0
    early_stop_count = 0
    while batch_num < config['max_its']:
        print('Batch Number: ',batch_num)

        #reset metrics
        train_loss.reset_states()
        train_acc_metric.reset_states()
        train_prec_metric.reset_states()
        train_rec_metric.reset_states()

        #Scores the training data and decides what to train on
        print('Getting Subset')
        train_DG.get_data_subset(model,train_ds)
        #wandb.log({'Train_loss_hist':wandb.Histogram(train_DG.losses)},step=batch_num)

        #Train on the data subset
        print('Training')
        for i in range(train_DG.num_batches):
            #get the activations for the next batch selection
            train_DG.get_activations(model,i)
            batch_data = train_DG.__getitem__(i)
            train_step(batch_data[0],batch_data[1])

        #Test on the test data
        print('Evaluating')
        test_metrics = model.evaluate(test_DG)

        #Log metrics
        wandb.log({'Train_loss':train_loss.result(),'Train_acc':train_acc_metric.result(),'Train_prec':train_prec_metric.result(),'Train_rec':train_rec_metric.result()},step=batch_num)
        wandb.log({'Test_loss':test_metrics[0],'Test_acc':test_metrics[1],'Test_prec':test_metrics[2],'Test_rec':test_metrics[3]},step=batch_num)
        batch_num += train_DG.num_batches

        if test_metrics[1] > early_stop_max:
            early_stop_max = test_metrics[1]
            early_stop_count = 0
        else:
            early_stop_count += train_DG.num_batches
            if early_stop_count > config['early_stop']:
                break

    wandb.log({'Images_used_hist':wandb.Histogram(train_DG.data_used)})
    #clear keras backend
    tf.keras.backend.clear_session()
    #Finish
    print('Finished')

if __name__ == "__main__":
    sweep_configuration = {
        'method': 'grid',
        'name': 'Test Acc vs Batch Size big batches 2',
        'metric': {
            'goal': 'maximize',
            'name': 'Test_acc'
            },
        'parameters': {
            'batch_size': {'values': [1024]},
            'learning_rate': {'values': [0.01,0.001,0.0001,0.00001]},
            'run_type': {'values': ['Random','SubMod']}, #,['Hard_Mining','SubMod'],['Hard_Mining','Random']
            'optimizer': {'values': ['SGD','Adam']},
            }
        }

    os.environ['WANDB_API_KEY'] = 'fc2ea89618ca0e1b85a71faee35950a78dd59744'
    wandb.login()
    sweep_id = wandb.sweep(sweep=sweep_configuration, project='k_diversity')
    wandb.agent(sweep_id, function=main)


from setuptools import find_packages, setup

base_dependencies = [
    "pandas>=1.0.0",
    "python-dateutil>=2",
]

additional_dependencies = {
    "dev": [
        "black>=21.9b0",
        "pre-commit>=2.15.0",
        "pytest>=6.2.1",
        "pylint>=2.7.4",
        "jupyterlab",
        "twine",
        "seaborn",
    ],
}

VERSION = "0.0.1"
DESCRIPTION = "R original utilities now in python"

with open("README.md", "r", encoding="utf8") as fh:
    LONG_DESCRIPTION = fh.read()

setup(
    name="rutil",
    packages=find_packages(where="rutil"),
    package_dir={"": "rutil"},
    description=DESCRIPTION,
    long_description=LONG_DESCRIPTION,
    long_description_content_type="text/markdown",
    author="Ismael Cabral",
    version="0.0.2",
    py_modules=["rutil"],
    keywords=["R language", "R", "str()", "pandas", "struture"],
    install_requires=base_dependencies,
    extras_require=additional_dependencies,
    python_requires=">=3.6",
)


from django.urls import path,re_path,include
# yazdigimiz functions imporr edirk
from .views import register,activate,login,logout

app_name="account"

urlpatterns = [
    path('register/', register, name='register'),
    re_path(r'^activate/(?P<uidb64>[0-9A-Za-z_\-]+)/(?P<token>[0-9A-Za-z]{1,13}-[0-9A-Za-z]{1,33})/$',activate, name='activate'),
    path('login/',login , name='login'),
    path('logout/',logout , name='logout'),

]


import itertools
from contextlib import suppress
from typing import Callable, Dict, Optional, Sequence, Union

import numpy as np
import pandas as pd

from .. import brightness, config, multicolor, spatial

try:
    import trackpy
    trackpy_available = True
except ImportError:
    trackpy_available = False

class SmFRETTracker:
    """Class for tracking of smFRET data

    There is support for dumping and loading to/from YAML using
    :py:mod:`sdt.io.yaml`.
    """
    yaml_tag = "!SmFRETTracker"
    _yaml_keys = ("excitation_seq", "registrator", "link_options",
                  "min_length", "brightness_options", "interpolate",
                  "coloc_dist", "acceptor_channel", "neighbor_radius")

    frame_selector: multicolor.FrameSelector
    """A :py:class:`FrameSelector` instance with the matching
    :py:attr:`excitation_seq`.
    """
    registrator: multicolor.Registrator
    """multicolor.Registrator used to overlay channels"""
    link_options: Dict
    """Options passed to :py:func:`trackpy.link_df`"""
    min_length: int
    """Minimum length of tracks"""
    neighbor_radius: float
    """How far two features may be apart while still being considered close
    enough so that one influences the brightness measurement of the other.
    This is related to the `radius` option of
    :py:func:`brightness.from_raw_image`.
    """
    brightness_options: Dict
    """Options passed to :py:func:`brightness.from_raw_image`. Make sure to
    adjust :py:attr:`neighbor_radius` if you change either the `mask` or the
    `radius` option!
    """
    interpolate: bool
    """Whether to interpolate coordinates of features that have been missed
    by the localization algorithm.
    """
    coloc_dist: float
    """After overlaying donor and acceptor channel features, this gives the
    maximum distance up to which donor and acceptor signal are considered
    to come from the same molecule.
    """
    acceptor_channel: int
    """Can be either 1 or 2, depending the acceptor is the first or the
    second channel in :py:attr:`registrator`.
    """
    columns: Dict
    """Column names in DataFrames. Defaults are taken from
    :py:attr:`config.columns`.
    """

    @config.set_columns
    def __init__(self, excitation_seq: str = "da",
                 registrator: Optional[multicolor.Registrator] = None,
                 link_radius: float = 5, link_mem: int = 1,
                 min_length: int = 1, feat_radius: int = 4,
                 bg_frame: int = 2,
                 bg_estimator: Union[str,
                                     Callable[[np.ndarray], float]] = "mean",
                 neighbor_radius: Union[float, str] = "auto",
                 interpolate: bool = True, coloc_dist: float = 2.0,
                 acceptor_channel: int = 2, link_quiet: bool = True,
                 link_options: Dict = {}, columns: Dict = {}):
        """Parameters
        ----------
        excitation_seq
            Set the :py:attr:`excitation_seq` attribute.
        registrator
            Registrator used to overlay channels. If `None`, use the identity
            transform.
        link_radius
            Maximum movement of features between frames. See `search_range`
            option of :py:func:`trackpy.link_df`.
        link_mem
            Maximum number of frames for which a feature may not be detected.
            See `memory` option of :py:func:`trackpy.link_df`.
        min_length
            Minimum length of tracks.
        feat_radius
            Radius of circle that is a little larger than features. See
            `radius` option of :py:func:`brightness.from_raw_image`.
        bg_frame
            Size of frame around features for background determination. See
            `bg_frame` option of :py:func:`brightness.from_raw_image`.
        bg_estimator
            Statistic to estimate background. See `bg_estimator` option of
            :py:func:`brightness.from_raw_image`.
        neighbor_radius
            How far two features may be apart while still being considered
            close enough so that one influences the brightness measurement of
            the other. This is related to the `radius` option of
            :py:func:`brightness.from_raw_image`. If "auto", use the smallest
            value that avoids overlaps.
        interpolate
            Whether to interpolate coordinates of features that have been
            missed by the localization algorithm.
        coloc_dist
            After overlaying donor and acceptor channel features, this gives
            the maximum distance up to which donor and acceptor signal are
            considered to come from the same molecule.
        acceptor_channel
            Whether the acceptor channel is number 1 or 2 in `registrator`.
        link_options
            Specify additional options to :py:func:`trackpy.link_df`.
            "search_range" and "memory" will be overwritten by the
            `link_radius` and `link_mem` parameters.
        link_quiet
            If `True`, call :py:func:`trackpy.quiet`.

        Other parameters
        ----------------
        columns
            Override default column names as defined in
            :py:attr:`config.columns`. Relevant names are `coords`, `time`,
            `mass`, `signal`, `bg`, `bg_dev`. This means, if your DataFrame has
            coordinate columns "x" and "z" and the time column "alt_frame", set
            ``columns={"coords": ["x", "z"], "time": "alt_frame"}``. This
            parameters sets the :py:attr:`columns` attribute.
        """
        self.frame_selector = multicolor.FrameSelector(excitation_seq)
        self.registrator = (registrator if registrator is not None
                            else multicolor.Registrator())

        self.link_options = link_options.copy()
        self.link_options["search_range"] = link_radius
        self.link_options["memory"] = link_mem

        self.min_length = min_length
        self.brightness_options = dict(
            radius=feat_radius,
            bg_frame=bg_frame,
            bg_estimator=bg_estimator,
            mask="circle")
        self.interpolate = interpolate
        self.coloc_dist = coloc_dist
        self.acceptor_channel = acceptor_channel
        self.columns = columns

        if isinstance(neighbor_radius, str):
            # auto radius
            neighbor_radius = 2 * feat_radius + 1
        self.neighbor_radius = neighbor_radius

        if link_quiet and trackpy_available:
            trackpy.quiet()

    @property
    def excitation_seq(self) -> str:
        """Excitation sequence. "d" stands for donor, "a" for acceptor,
        anything else describes other kinds of frames which are irrelevant for
        tracking.

        One needs only specify the shortest sequence that is repeated,
        i. e. "ddddaddddadddda" is the same as "dddda".
        """
        return self.frame_selector.excitation_seq

    @excitation_seq.setter
    def excitation_seq(self, seq: str):
        self.frame_selector.excitation_seq = seq

    def track(self, donor_img: Sequence[np.ndarray],
              acceptor_img: Sequence[np.ndarray], donor_loc: pd.DataFrame,
              acceptor_loc: pd.DataFrame, d_mass: bool = False
              ) -> pd.DataFrame:
        """Track smFRET data

        Localization data for both the donor and the acceptor channel is
        merged (since a FRET construct has to be visible in at least one
        channel). The merged data is than linked into trajectories using
        py:func:`trackpy.link_df`. For this the :py:mod:`trackpy` package needs
        to be installed. Additionally, the feature brightness is determined for
        both donor and acceptor for raw image data using
        :py:func:`brightness.from_raw_image`. These data are written into a
        a :py:class:`pandas.DataFrame` whose columns have a MultiIndex
        containing the "donor" and "acceptor" items in the top level.

        A column specifying whether the entry originates from donor or
        acceptor excitation is also added: ("fret", "exc_type"). It is "d"
        for donor and "a" for acceptor excitation; see the
        :py:meth:`flag_excitation_type` method.

        Parameters
        ----------
        donor_img, acceptor_img
            Raw image frames for donor and acceptor channel. This need to be
            of type `list`, but anything that returns image data when indexed
            with a frame number will do.
        donor_loc, acceptor_loc
            Localization data for donor and acceptor channel
        d_mass
            If `True`, get total brightness upon donor excitation by
            from the sum of donor and acceptor image. If `False`, the
            donor excitation brightness can still be calculated as the sum of
            donor and acceptor brightness in :py:meth:`analyze`. Defaults to
            `False`.

            Note: Until :py:mod:`slicerator` with support for multiple
            inputs to pipelines is released, setting this to `True` will load
            all of `donor_img` and `acceptor_img` into memory, even if
            :py:mod:`pims` is used.

        Returns
        -------
        The columns are indexed with a :py:class:`pandas.MultiIndex`.
        The top index level consists of "donor" (tracking data for the
        donor channel), "acceptor" (tracking data for the acceptor
        channel), and "fret". The latter contains a column with the
        particle number ("particle"), an indicator (0 / 1) whether there
        is a near neighbor ("has_neighbor"), and an indicator whether the
        data point was interpolated ("interp") because it was not in the
        localization data in either channel.
        """
        if not trackpy_available:
            raise RuntimeError("`trackpy` package required but not installed.")

        # Names of brightness-related columns
        br_cols = ["signal", "mass", "bg", "bg_dev"]
        # Position and frame columns
        posf_cols = self.columns["coords"] + [self.columns["time"]]

        # Don't modify originals
        donor_loc = donor_loc.copy()
        acceptor_loc = acceptor_loc.copy()
        for df in (donor_loc, acceptor_loc):
            for c in br_cols:
                # Make sure those columns exist
                df[c] = 0.

        donor_channel = 1 if self.acceptor_channel == 2 else 2
        donor_loc_corr = self.registrator(donor_loc, channel=donor_channel)

        # Create FRET tracks (in the acceptor channel)
        # Acceptor channel is used because in ALEX there are frames without
        # any donor locs, therefore minimizing the error by transforming
        # back and forth.
        coloc = multicolor.find_colocalizations(
                donor_loc_corr, acceptor_loc, max_dist=self.coloc_dist,
                channel_names=["donor", "acceptor"], keep_unmatched=True)
        coloc_pos_f = coloc.loc[:, (slice(None), posf_cols)]
        coloc_pos_f = coloc_pos_f.values.reshape(
            (len(coloc), 2, len(self.columns["coords"]) + 1))
        # Use the mean of positions as the new position
        merged = np.nanmean([coloc["donor"][posf_cols].values,
                             coloc["acceptor"][posf_cols].values], axis=0)
        merged = pd.DataFrame(merged, columns=posf_cols)
        merged["__trc_idx__"] = coloc.index  # works as long as index is unique

        self.link_options["pos_columns"] = self.columns["coords"]
        self.link_options["t_column"] = self.columns["time"]

        # Track only "d" and "a" frames. Renumber frames for that.
        merged = self.frame_selector.select(
            merged, "da", renumber=True, columns=self.columns)
        track_merged = trackpy.link_df(merged, **self.link_options)

        if self.interpolate:
            # Interpolate coordinates where no features were localized
            track_merged = spatial.interpolate_coords(track_merged,
                                                      self.columns)
        else:
            # Mark all as not interpolated
            track_merged["interp"] = 0

        # Flag localizations that are too close together
        if self.neighbor_radius:
            spatial.has_near_neighbor(track_merged, self.neighbor_radius,
                                      self.columns)

        # Restore original frame numbers that were changed when calling
        # selector.__call__ above.
        # Do this after calling spatial.interpolate_coords, otherwise the
        # excluded frames will be interpolated!
        self.frame_selector.restore_frame_numbers(track_merged, "da",
                                                  columns=self.columns)

        # Get non-interpolated colocalized features
        non_interp_mask = track_merged["interp"] == 0
        non_interp_idx = track_merged.loc[non_interp_mask, "__trc_idx__"]
        ret = coloc.loc[non_interp_idx]
        ret["fret", "particle"] = \
            track_merged.loc[non_interp_mask, "particle"].values

        # Append interpolated features (which "appear" only in the acceptor
        # channel)
        interp_mask = ~non_interp_mask
        interp = track_merged.loc[interp_mask, posf_cols]
        interp.columns = pd.MultiIndex.from_product([["acceptor"], posf_cols])
        interp["fret", "particle"] = \
            track_merged.loc[interp_mask, "particle"].values
        ret = pd.concat([ret, interp])

        # Add interp and has_neighbor column
        cols = ["interp"]
        if "has_neighbor" in track_merged.columns:
            cols.append("has_neighbor")
        ic = pd.concat([track_merged.loc[non_interp_mask, cols],
                        track_merged.loc[interp_mask, cols]])
        for c in cols:
            ret["fret", c] = ic[c].values

        # If coordinates or frame columns are NaN in any channel (which means
        # that it didn't have a colocalized partner), use the position and
        # frame number from the other channel.
        for c1, c2 in itertools.permutations(("donor", "acceptor")):
            mask = np.any(~np.isfinite(ret.loc[:, (c1, posf_cols)]), axis=1)
            d = ret.loc[mask, (c2, posf_cols)]
            ret.loc[mask, (c1, posf_cols)] = d.values

        # get feature brightness from raw image data
        ret_d = self.registrator(ret["donor"], channel=self.acceptor_channel)
        ret_a = ret["acceptor"].copy()
        brightness.from_raw_image(ret_d, donor_img, **self.brightness_options,
                                  columns=self.columns)
        brightness.from_raw_image(ret_a, acceptor_img,
                                  **self.brightness_options,
                                  columns=self.columns)

        # If desired, get brightness upon donor excitation from image overlay
        if d_mass:
            # Until slicerator with support for multiple inputs to pipelines
            # is released, load everything into memory
            overlay = []
            for di, da in zip(donor_img, acceptor_img):
                o = di + self.registrator(da, cval=np.mean,
                                          channel=self.acceptor_channel,
                                          columns=self.columns)
                overlay.append(o)
            df = ret_d[posf_cols].copy()
            brightness.from_raw_image(df, overlay, **self.brightness_options)
            ret["fret", "d_mass"] = df[self.columns["mass"]]

        ret_d.columns = pd.MultiIndex.from_product((["donor"], ret_d.columns))
        ret_a.columns = pd.MultiIndex.from_product((["acceptor"],
                                                    ret_a.columns))
        ret.drop(["donor", "acceptor"], axis=1, level=0, inplace=True)
        ret = pd.concat([ret_d, ret_a, ret], axis=1)
        ret.sort_values(
            [("fret", "particle"), ("donor", self.columns["time"])],
            inplace=True)

        # Filter short tracks (only count non-interpolated localizations)
        u, c = np.unique(
            ret.loc[ret["fret", "interp"] == 0, ("fret", "particle")],
            return_counts=True)
        valid = u[c >= self.min_length]
        ret = ret[ret["fret", "particle"].isin(valid)]

        ret = ret.reset_index(drop=True)
        self.flag_excitation_type(ret)
        return ret

    def flag_excitation_type(self, tracks: pd.DataFrame):
        """Add a column indicating excitation type (donor/acceptor/...)

        Add  ("fret", "exc_type") column in place. It is of "category" type.

        Parameters
        ----------
        tracks
            Result of :py:meth:`track`
        """
        # FIXME: Don't force convert to int, but raise an error (?)
        # First, the track method needs to preserve the data type of the time
        # column
        eseq = self.frame_selector.eval_seq()
        frames = tracks["donor", self.columns["time"]].to_numpy(dtype=int)
        et = pd.Series(eseq[frames % len(eseq)], dtype="category")
        # Assignment to dataframe is done by matching indices, not line-wise
        # Thus copy index
        et.index = tracks.index
        tracks["fret", "exc_type"] = et

    @classmethod
    def to_yaml(cls, dumper, data):
        """Dump as YAML

        Pass this as the `representer` parameter to
        :py:meth:`yaml.Dumper.add_representer`
        """
        m = {k: getattr(data, k) for k in cls._yaml_keys}
        return dumper.represent_mapping(cls.yaml_tag, m)

    @classmethod
    def from_yaml(cls, loader, node):
        """Construct from YAML

        Pass this as the `constructor` parameter to
        :py:meth:`yaml.Loader.add_constructor`
        """
        m = loader.construct_mapping(node)
        ret = cls()
        for k in cls._yaml_keys:
            setattr(ret, k, m[k])
        return ret

with suppress(ImportError):
    from ..io import yaml
    yaml.register_yaml_class(SmFRETTracker)


'''
An 'entity' is any object that populates the game (AKA, the content).
The 'Entity' class takes an x/y position, it's graphic/symbol, and color

Method:
    Entity.move(): Updates entity's x/y position with a new set of coordinates (called after a successful 'move' action)

'''

#_______________________________________________________________________// MODULES

from __future__ import annotations
import copy
from typing import (Optional, Tuple, Type, TypeVar, TYPE_CHECKING)

from render_order import RenderOrder

if TYPE_CHECKING:
    from components.ai import BaseAI
    from components.fighter import Fighter
    from game_map import GameMap

#_______________________________________________________________________// DECLARATION

# The 'Troll' enemy
T = TypeVar("T", bound="Entity")

#_______________________________________________________________________// CLASSES

# The generic object that represents players, enemies, items, etc.
class Entity:

    # Initialization
    def __init__(
        # Set initial values
        self,
        gamemap: Optional[GameMap] = None,
        x:                int = 0,                                    # x/y   - entity's position
        y:                int = 0,
        char:             str = "?",                                  # char  - its symbol/sprite
        color:            Tuple[int, int, int] = (255, 255, 255),     # color - an RGB color value
        name:             str = "<Unnamed>",                          # name  - references the entity
        blocks_movement:  bool = False,                               # Walkable or not (enemies aren't, items and NPCs are)
        render_order:     RenderOrder = RenderOrder.CORPSE            # 'corpse' renders first, so its a good default
    ):
        self.x              = x
        self.y              = y
        self.char           = char
        self.color          = color
        self.name           = name
        self.blocks_movement= blocks_movement
        self.render_order   = render_order

        if gamemap:
            # If a gamemap isn't provided now then it will be set later.
            self.gamemap = gamemap
            gamemap.entities.add(self)

    def spawn(self: T, gamemap: GameMap, x: int, y:int) -> T:
        '''
        Spawns a copy of this entity instance at a given location.
        '''
        clone = copy.deepcopy(self)

        clone.x = x
        clone.y = y

        clone.gamemap = gamemap
        gamemap.entities.add(clone)

        return clone

    def place(self, x: int, y: int, gamemap: Optional[GameMap] = None) -> None:
        '''
        Place this entity at a new location (handles moving between GameMaps).
        '''
        self.x = x
        self.y = y

        if gamemap:
            # If this entity has a GameMap associated with it:
            if hasattr(self, "gamemap"):
                self.gamemap.entites.remove(self)

            self.gamemap = gamemap
            gamemap.entities.add(self)

    def move(self, dx: int, dy: int) -> None:
        # Move the entity by a given amount
        self.x += dx
        self.y += dy

class Actor(Entity):
    '''
    Inherits all the properties/attributes/methods of the 'Entity' class.
    '''

    def __init__(
        self,
        *,
        x: int = 0,
        y: int = 0,
        char: str = "?",
        color: Tuple[int, int, int] = (255, 255, 255),
        name: str = "<Unnamed>",
        ai_cls: Type[BaseAI],
        fighter: Fighter
    ):
        # Calls the super class '__init__()' function (ie, the parent class / 'Entity' class __init__)
        super().__init__(
            x               = x,
            y               = y,
            char            = char,
            color           = color,
            name            = name,
            blocks_movement = True,                 # Always pass in 'True' (since all "actors" will block movement)
            render_order    = RenderOrder.ACTOR
        )

        # Assign components (each actor needs the ability to move around and deal/take damage)
        self.ai: Optional[BaseAI] = ai_cls(self)

        self.fighter = fighter
        self.fighter.entity = self

    @property
    def is_alive(self) -> bool:
        '''
        Returns 'True' as long as this actor has an 'ai'.
        - If the entity's '.die()' method gets called, its 'ai' is replaced with 'None', effectively removing its capacity to perform actions.
        - Calling 'die(). on this entity will also set this property to return 'False'.
        '''
        return bool(self.ai)


import cv2
import numpy as np
import pandas as pd
import os
import re

#同一品牌的两期图片链接数据判断，以图片链接为判断标准，分新上线和已下架

#获取文件名称
def get_filename(path):
    files = os.listdir(path)
    return files

#产品包装判断
def pic_sx_judge(filepath1,filepath2,name):
    df1 =  pd.read_excel(filepath1+name)
    df2 =  pd.read_excel(filepath2+name)
    df_same = pd.DataFrame(columns=df2.columns)
    df_diff = pd.DataFrame(columns=df2.columns)
    df_same['图片链接_old'] = np.NaN
    df_diff['图片链接_old'] = np.NaN
    df_same['页面网址_old'] = np.NaN
    df_diff['页面网址_old'] = np.NaN
    os.chdir(filepath4)

    for i in range(len(df2)):
        if re.search('.tmall',df2.loc[1,'页面网址']):

            url = re.search('id=(\d+)',df2.loc[i,'页面网址']).group(1)
            pic1 = df2.loc[i, '图片链接']
            for j in range(len(df1)):
                if url in df1.loc[j, '页面网址']:
                    pic2 = df1.loc[j, '图片链接']
                    try:
                        image1 = cv2.imread(pic1)
                        image2 = cv2.imread(pic2)
                        difference = cv2.subtract(image1, image2)
                        result = not np.any(difference)  # if difference is all zeros it will return False
                        if result is True:
                            df_same.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                            df_same.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                            df_same.loc[i, df2.columns] = df2.loc[i]

                        else:
                            df_diff.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                            df_diff.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                            df_diff.loc[i, df2.columns] = df2.loc[i]
                    except:
                        df_diff.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                        df_diff.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                        df_diff.loc[i, df2.columns] = df2.loc[i]
        else:
            url = df2.loc[i, '页面网址']
            pic1 = df2.loc[i, '图片链接']
            for j in range(len(df1)):
                if url in df1.loc[j, '页面网址']:
                    pic2 = df1.loc[j, '图片链接']
                    try:
                        image1 = cv2.imread(pic1)
                        image2 = cv2.imread(pic2)
                        difference = cv2.subtract(image1, image2)
                        result = not np.any(difference)  # if difference is all zeros it will return False
                        if result is True:
                            df_same.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                            df_same.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                            df_same.loc[i, df2.columns] = df2.loc[i]

                        else:
                            df_diff.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                            df_diff.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                            df_diff.loc[i, df2.columns] = df2.loc[i]
                    except:
                        df_diff.loc[i, '图片链接_old'] = df1.loc[j, '图片链接']
                        df_diff.loc[i, '页面网址_old'] = df1.loc[j, '页面网址']
                        df_diff.loc[i, df2.columns] = df2.loc[i]
    return (df_same,df_diff)

#商品上新判断
def new_judge(filepath1,filepath2):
    word = ['天猫/', '京东/']
    #word = ['京东/']
    for nm  in word:
        filename2 = get_filename(filepath2 + nm + '提取数据1/')
        print(filename2)
        for name in filename2:
            path1 = filepath1 + nm + '提取数据1/'
            path2 = filepath2 + nm + '提取数据1/'
            df = pic_sx_judge(path1, path2,name)
            if os.path.isfile(filepath2 + nm + '图片判断' + '/' + '相同图片/'+ name):
                os.remove(filepath2 + nm + '图片判断' + '/' + '相同图片/'+ name)
            if os.path.isfile(filepath2 + nm + '图片判断' + '/' + '不同图片/'+ name):
                os.remove(filepath2 + nm + '图片判断' + '/' + '不同图片/'+ name)
            df[0].to_excel(filepath2 + nm + '图片判断' + '/' + '相同图片/'+ name, index=False)
            df[1].to_excel(filepath2 + nm + '图片判断' + '/' + '不同图片/'+ name, index=False)
'''def match_pic_url(name1,name2,name3,num1=0, num2=1):
    df_right1 = pd.read_excel(name1)
    df_right2 = pd.read_excel(name2)
    df1 = pd.read_excel(name3,sheetname=num1)
    df2 = pd.read_excel(name3, sheetname=num2)
    writer = pd.ExcelWriter(name3)
    old_id = []
new_id = []
[old_id.append(re.search('id=(\d+)',df_right1.loc[i,'链接']).group(1))for i in range(len(df_right1))]
[new_id.append(re.search('id=(\d+)', df1.loc[i, '页面网址']).group(1)) for i in range(len(df1))]
    for i in range(len(df1)):
        url_id = re.search('id=(\d+)',df_right1.loc[i,'链接']).group(1)
        for j in range(len(df_right1)):
            if url_id in df_right1.loc[j,'链接']:
                df1.loc[i,'图片链接'] = df_right1.loc[j,'图片链接']
    for i in range(len(df2)):
        for j in range(len(df_right2)):
            if df2.loc[i,'页面网址'] == df_right2.loc[j,'链接']:
                df2.loc[i,'图片链接'] = df_right2.loc[j,'图片链接']
    if os.path.isfile(name3):
        os.remove(name3)
    df1.to_excel(writer, index=False, sheet_name='天猫')
    df2.to_excel(writer,index=False,sheet_name = '京东')
    writer.save()'''
def match_pic_url_single(name1,name3,name):
    df_right = pd.read_excel(name1)
    df = pd.read_excel(name3)
    writer = pd.ExcelWriter(name3)
    if len(df_right)>0:
        if re.search('tmall', df.loc[0, '页面网址']):
            for i in range(len(df)):
                print(df.loc[i, '页面网址'])
                id = re.search('id=(\d+)', df.loc[i, '页面网址']).group(1)
                for j in range(len(df_right)):
                    if id in df_right.loc[j,'链接'] :
                        df.loc[i, '图片链接'] = df_right.loc[j, '图片链接']
                        #print(df_right.loc[j, '图片链接'])
        else:
            for i in range(len(df)):
                for j in range(len(df_right)):
                    if df.loc[i, '页面网址'] == df_right.loc[j, '链接']:
                        df.loc[i, '图片链接'] = df_right.loc[j, '图片链接']
                        #print(df_right.loc[j, '图片链接'])
        if os.path.isfile(name3):
            os.remove(name3)
        df.to_excel(writer, index=False, sheet_name=name)
        writer.save()
def match_data(date):
    word = ['天猫/', '京东/']
    #word = ['京东/']
    for name  in word:
        filename1 = get_filename(filepath2+name+'提取数据1/')
        for name1 in filename1:
            print(name1)
            match_pic_url_single(filepath4 + date + name + name1,filepath2+ name + '提取数据1/' + name1,name.split('/')[0])
            #match_pic_url_single(filepath4 + '1.30/' + name + name1, filepath2 + name + '提取数据1/' + name1, name.split('/')[0])

# filepath1 = r'E:/01复硕正态/08项目/01沐浴露监测/02成品数据/旗舰店数据/2018-2-28/'
# filepath2 = r'E:/01复硕正态/08项目/01沐浴露监测/02成品数据/旗舰店数据/2018-3-31/'
# filepath4 = r'E:\A_judge_pic\2017-11-15/'

filepath1 = r'G:\01复硕正态\08项目\04中华项目监测\02成品数据\旗舰店数据\2018-3-31/'
filepath2 = r'G:\01复硕正态\08项目\04中华项目监测\02成品数据\旗舰店数据\2018-4-13/'
filepath4 = r'G:\A_judge_pic\toothpaste/'
if __name__ == '__main__':
    match_data('4.13/')
    #new_judge(filepath1, filepath2)


import numpy as np
import sympy
from sympy.solvers import solve
from sympy import Symbol
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import seaborn as sns
import argparse

def plot_figure2(df,
                 num_init_pts,
                 imbal,
                 taus,
                 trial,
                 sel_rnd,
                 num_classes,
                 dataset_name,
                 fig_dir):

    fig, ax = plt.subplots()

    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    df = (df[(df['sel_rnd']==sel_rnd) &
             (df['num_init_pts']==num_init_pts) &
             (df['imbal']==imbal) &
             (df['trial']==trial)][['tau',
                                    'DMGT_rare_amnt',
                                    'DMGT_common_amnt']])

    n = Symbol('n')

    ideal_sizes = np.asarray(list(map(lambda tau: int((num_classes/2)*np.ceil(solve(sympy.sqrt(n+1) - sympy.sqrt(n) - tau, n)[0])), taus)))

    plot = sns.lineplot(data=df,
                        x='tau',
                        y='DMGT_common_amnt',
                        ci=95,
                        color=sns.color_palette('pastel')[0],
                        estimator='mean',
                        linewidth=5)

    sns.lineplot(data=df,
                 x='tau',
                 y='DMGT_rare_amnt',
                 ci=95,
                 color=sns.color_palette('pastel')[3],
                 estimator='mean',
                 linewidth=5)

    sns.lineplot(x=taus,
                 y=ideal_sizes,
                 color='darkgrey',
                 linestyle='--',
                 linewidth=5)

    plt.tight_layout(rect=[0, 0, 1, 0.8])
    plot.set_xlabel('Threshold', fontsize=14)
    plot.set_ylabel('Size of Selected Set', fontsize=14)
    plot.legend(labels=['common','rare', 'ideal'], bbox_to_anchor=(0.7,0.3), prop={'size':8})
    fig.savefig(fig_dir + f'{dataset_name}_tau_v_size.pdf', bbox_inches = "tight")


import math
import pandas as pd
import numpy as np

import argparse

import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D

import matplotlib.pyplot as plt
from torch import norm

def main(args):
    df = pd.read_csv(args.filename)
    df["No"] = range(0, len(df.index))

    # calc K constant
    df["K"] = (
        np.sqrt(df["velocity_x"] ** 2 + df["velocity_y"] ** 2)
        * df["integration_time"]
        / (df["range"] * np.sqrt(df["delta_px"] ** 2 + df["delta_py"] ** 2))
    )

    # calc velocity magnitude
    df["velocity"] = np.sqrt(df["velocity_x"] ** 2 + df["velocity_y"] ** 2)

    # calc delta magnitude
    df["delta"] = np.sqrt(df["delta_px"] ** 2 + df["delta_py"] ** 2)

    # remove NaN and inf
    df = df.replace([np.inf, -np.inf], np.nan).dropna()

    # remove zero velocity
    df = df[df["velocity"] != 0]

    # remove zero delta
    df = df[df["delta"] != 0]

    # add histogram of K constant
    plt.hist(
        df[df["surface_quality"] > 230]["K"], bins=100, alpha=0.5, color="red", log=True
    )

    k_hist, k_bin_edges = np.histogram(df[df["surface_quality"] > 230]["K"], bins=100)
    k_classvalues = (k_bin_edges[:-1] + k_bin_edges[1:]) / 2

    k_list = np.array(list(zip(k_classvalues, k_hist)))
    # sort list by frequency
    k_list = k_list[k_list[:, 1].argsort()[::-1]]

    # print top 5 K constant
    print(f"Top 5 K constant: {k_list[:5]}")

    # determine K constant, use top 5 K constant, weighted by frequency
    k_constant = np.average(k_list[:5, 0], weights=k_list[:5, 1])
    print(f"K constant: {k_constant}")

    """
    # print mode
    k_modes = df[df["surface_quality"] > 230]["K"].mode()
    print(f"K modes: {k_modes}")

    # remove 0 from mode
    k_modes = k_modes[k_modes != 0]

    # determine K constant, use top 3 modes
    k_constant = k_modes.head(1).mean()
    print(f"K constant: {k_constant}")
    """

    plt.show()

    # show K constant
    plt.scatter(
        df[df["surface_quality"] > 230]["No"], df[df["surface_quality"] > 230]["K"]
    )
    plt.axhline(y=k_constant, color="r", linestyle="-")
    plt.show()

    # separate data to 16 bins based on surface quality
    # surface quality is an integer from 0 to 255
    # 255 is the best quality, 0 is the worst
    # 16 bins are created
    df_bins = []
    for i in range(16):
        df_bins.append(
            df[
                (df["surface_quality"] >= (i * 16))
                & (df["surface_quality"] < ((i + 1) * 16))
            ]
        )

    centerized_df_bins = []

    # centerize delta_px and delta_py
    for df_bin in df_bins:
        centerized_df_bins.append(
            df_bin.assign(
                delta_px=lambda x: x["delta_px"]
                - x["velocity_x"] * x["integration_time"] / x["range"] / k_constant,
                delta_py=lambda x: x["delta_py"]
                - x["velocity_y"] * x["integration_time"] / x["range"] / k_constant,
                delta=lambda x: x["delta"]
                - x["velocity"] * x["integration_time"] / x["range"] / k_constant,
                actual_delta_px=lambda x: x["velocity_x"]
                * x["integration_time"]
                / x["range"]
                / k_constant,
                actual_delta_py=lambda x: x["velocity_y"]
                * x["integration_time"]
                / x["range"]
                / k_constant,
            )
        )

        # drop NaN and inf
        centerized_df_bins[-1] = (
            centerized_df_bins[-1].replace([np.inf, -np.inf], np.nan).dropna()
        )

    # calculate covariance for each bin
    covs = []
    for df_bin in centerized_df_bins:
        covs.append(np.cov(df_bin["delta_px"], df_bin["delta_py"]))

    # plot 16 bins
    # add scatter plot of delta_px and delta_py, and histogram of delta_px and delta_py
    hist_axs = []
    fig, axs = plt.subplots(4, 4, figsize=(20, 20))
    for i in range(16):
        axs[i // 4, i % 4].scatter(
            centerized_df_bins[i]["delta_px"], centerized_df_bins[i]["delta_py"]
        )
        axs[i // 4, i % 4].set_title(
            f"surface_quality: {i*16} - {(i+1)*16}, \n cov: {covs[i]}, \n actual dpx: {centerized_df_bins[i]['actual_delta_px'].mean()}"
        )
        axs[i // 4, i % 4].set_xlabel("delta_px")
        axs[i // 4, i % 4].set_ylabel("delta_py")
        axs[i // 4, i % 4].set_xlim([-1, 1])
        axs[i // 4, i % 4].set_ylim([-1, 1])
        axs[i // 4, i % 4].grid()

        hist_axs.append(axs[i // 4, i % 4].twinx())

        hist_axs[i].hist(
            centerized_df_bins[i]["delta_px"], bins=20, alpha=0.5, color="red"
        )
        hist_axs[i].hist(
            centerized_df_bins[i]["delta_py"], bins=20, alpha=0.5, color="blue"
        )

    # add space between plots
    fig.tight_layout(pad=5.0)

    # show plot
    plt.show()

    fig, axs = plt.subplots(4, 4, figsize=(20, 20), subplot_kw={"projection": "3d"})

    # plot 3d histogram of delta, velocity, separated by surface quality
    for i in range(4):
        for j in range(4):
            ax = axs[i, j]

            # Create histogram bins
            delta_bins = np.linspace(-1.0, 1.0, 20)
            velocity_bins = np.linspace(0, 0.2, 20)

            df_bin = centerized_df_bins[i * 4 + j]

            # Create histogram
            hist, xedges, yedges = np.histogram2d(
                df_bin["delta"],
                df_bin["velocity"],
                bins=(delta_bins, velocity_bins),
            )
            # Normalize each velocity bin
            for k in range(hist.shape[1]):
                total = np.sum(hist[:, k])
                if total > 0:
                    hist[:, k] /= total

            # Construct arrays for the anchor positions of the bars.
            xpos, ypos = np.meshgrid(xedges[:-1], yedges[:-1], indexing="ij")
            xpos = xpos.flatten()
            ypos = ypos.flatten()
            zpos = 0

            # Construct arrays with the dimensions for the bars.
            dx = dy = 0.1
            dz = hist.flatten()

            ax.bar3d(xpos, ypos, zpos, dx, dy, dz, shade=True)
            ax.set_xlabel("delta")
            ax.set_ylabel("velocity")
            ax.set_zlabel("Frequency")

            # set z range
            ax.set_zlim([0, 1.0])

    plt.show()

    # plot 2d histogram of delta, velocity, separated by surface quality
    fig, axs = plt.subplots(4, 4, figsize=(20, 20))
    for i in range(4):
        for j in range(4):
            ax = axs[i, j]
            ax.set_title(
                f"surface_quality: {(i*4+j)*16} - {(i*4+j+1)*16}, \n cov: {covs[i*4+j]}, \n actual dpx: {centerized_df_bins[i*4+j]['actual_delta_px'].mean()}"
            )

            df_bin = centerized_df_bins[i * 4 + j]

            ax.hist2d(
                df_bin["delta"],
                df_bin["velocity"],
                bins=[25, 10],
                cmap="Blues",
                range=[[-1, 1], [0, 0.2]],
            )
            ax.set_xlabel("delta")
            ax.set_ylabel("velocity")
    fig.tight_layout(pad=5.0)
    plt.show()

if __name__ == "__main__":
    # get filename from command line
    parser = argparse.ArgumentParser(
        description="Calculate covariance of optical flow measurements."
    )
    parser.add_argument(
        "filename", type=str, help="CSV file containing optical flow measurements."
    )
    args = parser.parse_args()

    main(args)


# Create your views here.
from rest_framework import mixins, status
from rest_framework import viewsets
from rest_framework import permissions
from rest_framework.request import Request
from rest_framework.response import Response
from rest_framework import mixins
from rest_framework.views import APIView

from account.models import Account
from account.schemas import account_list
from account.serializers import AccountSerializer, CustomTokenObtainPairSerializer, CustomRefreshTokenSerializer, \
    RegisterSerializer
from rest_framework_simplejwt import views
from rest_framework_simplejwt.serializers import TokenObtainPairSerializer

from back import settings

class RegisterApiView(mixins.CreateModelMixin, viewsets.GenericViewSet):
    serializer_class = RegisterSerializer
    queryset = Account.objects.all()

    def create(self, request, *args, **kwargs):
        serializer = RegisterSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data, status=status.HTTP_201_CREATED)
        return Response(serializer.errors, status=status.HTTP_409_CONFLICT)
    # def pe

class LogoutApiView(APIView):

    def post(self, request, *args, **kwargs):
        rt = request.COOKIES.get('refresh', None)
        if rt is None:
            return Response({"detail": "Refresh token not found"}, status=400)
        response = Response()
        response.delete_cookie('refresh')
        response.delete_cookie('access')
        return response

class AccountViewSet(mixins.ListModelMixin, viewsets.GenericViewSet):
    permission_classes = [permissions.IsAuthenticated]

    @account_list
    def list(self, request: Request, *args, **kwargs):
        user_id = request.query_params.get('user_id')
        queryset = Account.objects.get(id=user_id)
        serializer = AccountSerializer(queryset)
        return Response(serializer.data)

class JWTCookieMixin(APIView):
    def finalize_response(self, request: Request, response: Response, *args, **kwargs):
        if response.data.get('access'):
            response.set_cookie('access',
                                response.data['access'],
                                max_age=settings.SIMPLE_JWT["ACCESS_TOKEN_LIFETIME"],
                                httponly=True,
                                samesite='Strict')
            del response.data['access']

        if response.data.get('refresh'):
            response.set_cookie('refresh',
                                response.data['refresh'],
                                max_age=settings.SIMPLE_JWT["REFRESH_TOKEN_LIFETIME"],
                                httponly=True,
                                samesite='Strict')
            del response.data['refresh']

        return super().finalize_response(request, response, *args, **kwargs)

class JWTCookieTokenPairView(JWTCookieMixin, views.TokenObtainPairView):
    serializer_class = CustomTokenObtainPairSerializer

class JWTCookieRefreshView(JWTCookieMixin, views.TokenRefreshView):
    serializer_class = CustomRefreshTokenSerializer


from unittest import TestCase
import datetime
from entities import patient, doctor
from services import service_gateway, entity_service, doctors_service, patients_service

# Tests for patients service API
class PatientServiceTest(TestCase):
    def test_patientAPI(self):
        es = entity_service.EntityService(entity_service.DictionaryBackingStore())
        service_gateway.Registry.registerService(es)
        ds = service_gateway.Registry.lookupService('DoctorsService')
        ps = service_gateway.Registry.lookupService('PatientsService')

        d = doctor.Doctor('Dolittle', datetime.date(1920, 5, 1), 'M', 'MD')
        p = patient.Patient('Buratino', datetime.date(2010, 1, 1), 'M', 60, 30, d.Id)
        p.Preconditions = ['Made of wood']
        p.Allergies = ['Water']
        p.Medications = ['Read a book twice a day']

        dctx = service_gateway.InvocationContext(d.Id)
        pctx = service_gateway.InvocationContext(p.Id)
        ds.createPatient(dctx, p)

        p1 = ds.findPatient(dctx, p.Id)
        p2 = ps.findPatient(pctx, p.Name)
        self.assertEqual(p1, p2)

        p2.Allergies.append('Milk')
        ps.updatePatient(pctx, p2)

        p3 = ps.findPatient(pctx, p.Name)
        self.assertIn('Milk', p3.Allergies)


from schema import Schema, And, Or, Optional, Use
import copy
from .base import base_mod, TARGETS, ON_EVENTS, _on_effect, _buff, _constraints, _from, _creates

_cwl = {
  Optional("from"): _from, # consumed when the event triggers, if any
  Optional("on"): Or( # not _really_ optional but High Energy Fire doesn't have an `on` so ... ?
    Or(*ON_EVENTS),
    _on_effect,
    [_on_effect]
  ),
  Optional("target"): Or(*TARGETS), # what the effect targets
  Optional("stacks"): bool, # can I have > 1 of this mod
  Optional("buffs"): _buff, # does this create a buff?
  Optional("constraint"): _constraints,

}
cwl_creates = copy.copy(base_mod)
cwl_creates.update(_cwl)
cwl_uses = copy.copy(cwl_creates)

cwl_creates.update({
  "creates": _creates,
  "target": Or(*TARGETS),
})

cwl_uses.update({
  Optional("creates"): _creates
})

cwl_creates = Schema(cwl_creates)
cwl_uses = Schema(cwl_uses)


#%%  Imports
import os.path as osp
from skimage import io
from skimage.transform import resize
import numpy as np
import pandas as pd
from tqdm import tqdm
from PIL import Image, ImageDraw, ImageFont
from random import shuffle, seed

#%%  Calculate sizes and resolutions

path = r'c:\dev\a2_contact'
images_grid = ( 4, 8 )
seeds = ( 448, 444 )
im_dim = ( 420, 320 )
border_ratio = ( .8, .5 )

a4_mm = ( 210, 297 )
full_image_pixels = tuple( 10 * x for x in a4_mm )

im_diff = ( 40, 40 )
im_border = [0,0]

for i in [0,1]:
    im_border[i] = int( border_ratio[i] * (full_image_pixels[i] - ( im_dim[i] * images_grid[i] + im_diff[i] * (images_grid[i] - 1) ))  )

#%%  Some font issues calculations
chars = 1488 + np.arange( 27 )
char_length = ( 20, 50, 90, 50, 20, 100, 80, 20, 20, 100, 50, 50, 60, 10, 20, 100, 100, 20, 60, 50, 60, 60, 40, 20, 50, 10, 20 )
clm = dict( zip( chars, map( lambda x:(190 - x)/10, char_length )) )
chars = ord( '0' ) + np.arange( 10 )
clm.update( zip( chars, 10 * [ 13 ] ) )
clm[ ord( ' ' ) ] = 7
clm[ ord( '.' ) ] = 7
clm[ ord( '-' ) ] = 10

#%%  Read from excel
df = pd.ExcelFile( osp.join( path, r'contacts.xlsx' ) ).parse( 'Sheet2' )
contacts = []
contact_items = ( 'image', 'name', 'date', 'ima_name', 'ima_num', 'aba_name', 'aba_num' )
for i, row in df.iterrows():
    contacts.append( dict( zip( contact_items, [ row[ x ] for x in contact_items ] ) ) )

#%%  Reading contact images
for c in tqdm( contacts ):
    c[ 'im_data' ] = resize( io.imread( osp.join( path, c['image'] ) ), ( im_dim[1], im_dim[1] ) )

#%%  shuffle contacts
for s in seeds:
    seed( s )
    shuffle( contacts )

#%%  generating text image
def GetTxtAndMaxLength( c ):
    lines = [ c['date'] + ' ' + c['name'][-1::-1],
              c['ima_num'] + ' ' + c['ima_name'][-1::-1] ]
    if isinstance( c['aba_name'], str ):
        lines.append( c['aba_num'] + ' ' + c['aba_name'][-1::-1] )
    max_length = np.max( [ sum( [ clm[ord(x)] for x in txt ] ) for txt in lines ] )
    return '\n'.join( lines ), max_length

def EmbedTxtToImg( im, c ):
    pil_im = Image.fromarray( ( 255 * im ).astype( np.uint8 ) )
    fnt = ImageFont.truetype( font='david.ttf', size=30, index=0, encoding='unic', layout_engine=ImageFont.LAYOUT_RAQM )
    txt, max_length = GetTxtAndMaxLength( c )
    draw = ImageDraw.Draw( pil_im )
    draw.multiline_text( ( im.shape[1] - max_length - 4, 4 ), txt, fill=(255,255,0), font=fnt, spacing=4, align="right" )
    return np.array( pil_im ).astype( np.float ) / 255

#%% create the final image
full_im = resize( io.imread( osp.join( path, r'school-rocket-ship.jpeg' )), ( full_image_pixels[0], full_image_pixels[1] ) )
for (y,x), c in tqdm( zip( np.ndindex( images_grid ), contacts ) ):
    y_start = y * (im_dim[0] + im_diff[0]) + im_border[0]
    y_end = y_start + im_dim[1]
    x_start = x * (im_dim[1] + im_diff[1]) + im_border[1]
    x_end = x_start + im_dim[1]
    full_im[ y_start : y_end, x_start : x_end ] = contacts[ y * images_grid[1] + x ][ 'im_data' ]
    y_text_start = y_end
    y_text_end = y_text_start + im_dim[0] - im_dim[1]
    full_im[ y_text_start : y_text_end, x_start : x_end ] = EmbedTxtToImg( .4 * full_im[ y_text_start : y_text_end, x_start : x_end ], c )

#%% save to file
io.imsave( osp.join( path, 'contacts.png' ), full_im )


import argparse
import time
import logging
import json
import csv
import typing
from datetime import datetime
from constant import tables, table_schema2, table_schema
import apache_beam as beam
from apache_beam.io import fileio
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.transforms.trigger import AfterWatermark, AfterCount, AfterProcessingTime, AccumulationMode
from apache_beam.transforms.combiners import CountCombineFn
from apache_beam.runners import DataflowRunner, DirectRunner
#from IPython import embed

# Table schema for BigQuery
"""table_schema = {
    'fields':[
        {
            "name": "name",
            "mode": "NULLABLE",
            "type": "STRING"
        },
        {
            "name": "age",
            "mode": "NULLABLE",
            "type": "INTEGER"
        },
        {
            "name": "height",
            "mode": "NULLABLE",
            "type": "FLOAT"
        }
    ]
}"""
# ### functions and classes

def getSchema(element):
    names=element["address"].split("/")
    if not names[-1] in tables:
        return tables["tabledata.csv"]
    return tables[names[-1]]

def parse_json(element):
    row = json.loads(element.decode('utf-8'))
    #embed()
    return row #CommonLog(**row)

def parse_csv(element, message):
    lines=element.split(",")
    #head=lines[0]
    #lines.pop(0)
    schema=tables["tabledata.csv"]["fields"]
    names=message["address"].split("/")
    if names[-1] in tables:
        schema= tables[names[-1]]["fields"]

    #schema=getSchema(message)["fields"]
    dict={}
    for col in range(len(schema)):
        dict[schema[col]["name"]]=lines[col]
    #embed()
    return dict

def getURI(element):
    return element["address"]

def getDest(element):
    return element["destination"]

def cross_join(left, rights):
    #embed()
    return [left, rights]

# ### main

def run():
    # Command line arguments
    parser = argparse.ArgumentParser(description='Load from Json from Pub/Sub into a dynamic BigQuery output')

    # Setting up the Beam pipeline options
    options = PipelineOptions( save_main_session=True, streaming=True)
    options.view_as(GoogleCloudOptions).project = 'smartlive'
    options.view_as(GoogleCloudOptions).region = 'europe-west1'
    options.view_as(GoogleCloudOptions).staging_location = 'gs://testinsertbigquery/staging'
    options.view_as(GoogleCloudOptions).temp_location = 'gs://testinsertbigquery/temp'
    options.view_as(GoogleCloudOptions).dataflow_service_options=["enable_prime"]
    options.view_as(DebugOptions).experiments=["use_runner_v2"]
    options.view_as(GoogleCloudOptions).job_name = '{0}{1}'.format('chargement-fichier-streaming-primerunnerv2',time.time_ns())
    options.view_as(StandardOptions).runner = 'DataflowRunner'

    input_topic = 'projects/smartlive/topics/my_topic2'

    dataset='EtienneResults'

    table_names = ['smartlive:EtienneResults.output1', 'smartlive:EtienneResults.output2', 'smartlive:EtienneResults.output3']

    def table_fn(row, element):
        #embed()
        destination=element["destination"]
        return options.view_as(GoogleCloudOptions).project + ":" + dataset + "."+ destination

    window_duration = 60
    ##allowed_lateness = opts.allowed_lateness
    ##dead_letter_bucket = opts.dead_letter_bucket

    p = beam.Pipeline(options=options)

    readData=(p
            | 'ReadFromPubSub' >> beam.io.ReadFromPubSub(input_topic)
            | 'ParseJson' >> beam.Map(parse_json))

    dest_view=beam.pvalue.AsSingleton(
        readData | "Window1 into Fixed Intervals" >> beam.WindowInto
        (beam.window.FixedWindows(1),
         trigger=AfterCount(1),
         accumulation_mode=AccumulationMode.DISCARDING))

    data_view=(readData
            | 'Get CSV address' >> beam.Map(getURI)
            | 'Read CSV files' >> beam.io.ReadAllFromText(skip_header_lines=1)
            #| beam.GroupBy(lambda s: 0)
            #| 'ParseCSV' >> beam.Map(parse_csv)
            | "Window2 into Fixed Intervals" >> beam.WindowInto
        (beam.window.FixedWindows(1),
         trigger=AfterCount(1),
         accumulation_mode=AccumulationMode.DISCARDING))

    (data_view
        | 'ParseCSV' >> beam.Map(parse_csv, message=dest_view)
        | 'WriteWithDynamicDestination' >> beam.io.WriteToBigQuery(
                table=lambda row, dest_view : table_fn(row,dest_view),
                schema=lambda row, dest_view : getSchema(dest_view),
                table_side_inputs=(dest_view, ),
                schema_side_inputs=(dest_view, ),
                #additional_bq_parameters={'ignoreUnknownValues': True},
                create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
                write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND
                ))

    logging.getLogger().setLevel(logging.INFO)
    logging.info("Building pipeline ...")

    p.run().wait_until_finish()

if __name__ == '__main__':
  run()


"""
HTTP methods
"""
import json
import urllib
import requests

from .. import env
from ..exceptions.bitpay_exception import BitPayException
from .key_utils import get_compressed_public_key_from_pem, sign

class RESTcli:
    __headers = {}
    __baseurl = ""
    __eckey = ""
    __identity = ""
    __proxy = ""

    def __init__(self, environment, eckey, proxy=None):
        self.__eckey = eckey
        self.__baseurl = env.TESTURL if environment == env.TEST else env.PRODURL
        self.__proxy = proxy
        self.init()

    def init(self):
        try:
            self.__headers = {
                "x-accept-version": env.BITPAYAPIVERSION,
                "x-bitpay-plugin-info": env.BITPAYPLUGININFO,
                "x-bitpay-api-frame": env.BITPAYAPIFRAME,
                "x-bitpay-api-frame-version": env.BITPAYAPIFRAMEVERSION,
                "content-type": "application/json",
                "X-accept-version": "2.0.0"
            }
            if self.__proxy != "":
                self.__headers["proxy"] = self.__proxy

        except BitPayException as e:
            print(e)

    def post(self, uri, form_data, signature_required=True):
        """
        :param uri: Url at which user wants to post the data
        :param form_data: the data to be posted
        :param signature_required: Signature of the full request URL concatenated
        with the request body
        :return: json response
        """
        full_url = self.__baseurl + uri
        form_data = json.dumps(form_data)
        if signature_required:
            self.__headers["x-signature"] = sign(full_url + form_data, self.__eckey)
            self.__headers["x-identity"] = get_compressed_public_key_from_pem(self.__eckey)

        response = requests.post(full_url, data=form_data, headers=self.__headers)
        json_response = self.response_to_json_string(response)
        return json_response

    def get(self, uri, parameters=None, signature_required=True):
        """

        :param uri: Url from which user wants to get the data
        :param parameters: These are query parameters
        :param signature_required: Signature of the full request URL concatenated
        :return: json response
        """
        full_url = self.__baseurl + uri
        if parameters is not None:
            full_url = "%s?%s" % (full_url, urllib.parse.urlencode(parameters))

        if signature_required:
            self.__headers["x-signature"] = sign(full_url, self.__eckey)
            self.__headers["x-identity"] = get_compressed_public_key_from_pem(self.__eckey)

        response = requests.get(full_url, headers=self.__headers)
        json_response = self.response_to_json_string(response)
        return json_response

    def delete(self, uri, parameters=None):
        """

        :param uri: Url from which user wants to delete the data
        :param parameters: These are query parameters
        :return: json response
        """
        full_url = self.__baseurl + uri

        if parameters is not None:
            full_url = "%s?%s" % (full_url, urllib.parse.urlencode(parameters))

        self.__headers["x-signature"] = sign(full_url, self.__eckey)
        self.__headers["x-identity"] = get_compressed_public_key_from_pem(self.__eckey)

        response = requests.delete(full_url, headers=self.__headers)
        json_response = self.response_to_json_string(response)
        return json_response

    def update(self, uri, form_data):
        """
        :param uri: Url from which user wants to delete the data
        :param form_data: the data to be updated
        :return: json response
        """
        full_url = self.__baseurl + uri
        form_data = json.dumps(form_data)

        self.__headers["x-signature"] = sign(full_url + form_data, self.__eckey)
        self.__headers["x-identity"] = get_compressed_public_key_from_pem(self.__eckey)

        response = requests.put(full_url, data=form_data, headers=self.__headers)
        json_response = self.response_to_json_string(response)
        return json_response

    def response_to_json_string(self, response):
        if not response.json():
            raise BitPayException("Error: HTTP response is null")

        response_obj = response.json()
        if "status" in response_obj:
            if response_obj["status"] == "error":
                raise BitPayException(
                    "Error: " + response_obj["error"], response_obj["code"]
                )

        if "error" in response_obj:
            raise BitPayException(
                "Error: " + response_obj["error"]
            )
        elif "errors" in response_obj:
            message = ""
            for error in response_obj["errors"]:
                message += "\n" + str(error)
            raise BitPayException("Errors: " + message)

        if "success" in response_obj:
            return response_obj["success"]

        if "data" in response_obj:
            return response_obj["data"]

        return response_obj


from django.conf.urls import url
from . import views

urlpatterns = [
    url(r"login$",views.login),
    url(r"login_handler$",views.login_handler),
    url(r"register$",views.register),
    url(r"register_handler$",views.register_handler),
    url(r"register_exist/$",views.register_exist),
    url(r'user_info$',views.user_info),
    url(r'user_order/(.*)?$',views.user_order),
    url(r'user_site$',views.user_site),
    url(r"logout$",views.logout)

]


import argparse
from pathlib import Path
import os
import torch
from torch import nn
from torch.utils.data import DataLoader

from experiments.exp_def import TaskDefs, EncoderModelType
from data_utils.task_def import TaskType
from torch.utils.data import DataLoader
from mt_dnn.batcher import SingleTaskDataset, Collater
from mt_dnn.model import MTDNNModel
from mt_dnn.inference import eval_model
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import pyplot as plt
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available(),
                    help='whether to use GPU acceleration.')
parser.add_argument('--device_id', type=int, default=0)
parser.add_argument('--task', type=str)
parser.add_argument('--subtask', type=str, default='head_probe')

args = parser.parse_args()
TASK = args.task
SUBTASK = args.subtask
MAX_SEQ_LEN = 512
BATCH_SIZE = 8

if TASK == 'NLI':
    task_type = TaskType.Classification
elif TASK == 'POS':
    task_type = TaskType.SequenceLabeling
elif TASK == 'NER':
    task_type = TaskType.Classification
else:
    raise NotImplementedError

def save_acc_matrix_as_heatmap(input_csv='', file_df=None, invert_y=False, axlabels=None, out_file=''):
    if file_df is None:
        file_df = pd.read_csv(input_csv, index_col=0)
        assert len(out_file) > 0

    row_labels = file_df.index
    column_labels = file_df.columns
    data = file_df.to_numpy()

    plt.figure(figsize=(14, 14))
    annot_kws = {
        "fontsize":14,
    }
    heatmap = sns.heatmap(
        data,
        cbar=False,
        annot=True,
        annot_kws=annot_kws,
        fmt=".2f",
        # xticklabels=column_labels,
        # yticklabels=row_labels,
        cmap='RdYlGn')

    if invert_y:
        heatmap.invert_yaxis()
    if axlabels is not None:
        heatmap.set_ylabel(axlabels[1], fontsize=14)
        heatmap.set_xlabel(axlabels[0], fontsize=14)

    heatmap.set_yticklabels(row_labels, rotation=0, fontsize=14)
    heatmap.set_xticklabels(column_labels, rotation=0, fontsize=14)
    fig = heatmap.get_figure()

    if len(input_csv) > 0:
        fig.savefig(Path(input_csv).with_suffix('.png'))
    else:
        fig.savefig(Path(out_file).with_suffix('.pdf'), bbox_inches='tight')

def build_dataset(data_path, encoder_type, batch_size, max_seq_len, task_def, device_id):
    test_data_set = SingleTaskDataset(
        path=data_path,
        is_train=False,
        maxlen=max_seq_len,
        task_id=0,
        task_def=task_def
    )

    collater = Collater(is_train=False, encoder_type=encoder_type)

    test_data = DataLoader(
        test_data_set,
        batch_size=batch_size,
        collate_fn=collater.collate_fn,
        pin_memory=device_id > 0
    )

    return test_data

def get_acc(model, test_data, metric_meta, device_id, head_probe):
    with torch.no_grad():
        model.network.eval()
        model.network.to(device_id)

        results = eval_model(
            model,
            test_data,
            task_type=task_type,
            metric_meta=metric_meta,
            device=device_id,
            with_label=True,
            head_probe=head_probe
        )
    metrics = results[0]
    predictions = results[1]
    golds = results[3]
    ids = results[4]

    preds_df = pd.Series(predictions)
    golds_df = pd.Series(golds)
    id_df = pd.Series(ids)
    return metrics['ACC'], preds_df, golds_df, id_df

def evaluate_model_against_multiple_datasets(model,
                                            datasets,
                                            task_def,
                                            encoder_type,
                                            batch_size,
                                            max_seq_len,
                                            device_id):
    assert TASK == 'NLI'
    accs = []

    for dataset in datasets:
        print(f'\tevaluating on {dataset}')
        data_path = f'experiments/attention-probing/{dataset}/bert-base-multilingual-cased/{dataset}_test.json'
        test_data = build_dataset(data_path, encoder_type, batch_size, max_seq_len, task_def, device_id)
        acc = get_acc(model, test_data, metric_meta, device_id, head_probe=False)
        accs.append(acc)

    return accs

def evaluate_head_probe(hlhis, device_id, task_def, encoder_type, batch_size, max_seq_len, root_ckpt_path, setting, data_path, task_name, return_base_model):
    assert setting in ['multi', 'cross', 'base']

    if setting == 'cross':
        shorthand = 'cl'
    elif setting == 'multi':
        shorthand = 'ml'
    elif setting == 'base':
        shorthand = 'base'

    test_data = build_dataset(
        data_path,
        encoder_type,
        batch_size,
        max_seq_len,
        TaskDefs(task_def).get_task_def(task_name),
        device_id)

    if not return_base_model:
        base_model_ckpt = root_ckpt_path.joinpath(f'{setting}/model_5.pt')
    else:
        # just for config stuff to set up MTDNN.
        base_model_ckpt = root_ckpt_path.joinpath(f'cross/model_5.pt')

    model, metric_meta = construct_model(base_model_ckpt, task_name, task_def, return_base_model=return_base_model)

    if TASK == 'NLI':
        n_classes = 3
    elif TASK == 'POS':
        n_classes = 17
    elif TASK == 'NER':
        n_classes = 7

    output_dir = Path(f'score_outputs/{TASK}/{SUBTASK}/{setting}')
    output_dir.mkdir(parents=True, exist_ok=True)

    for i, (hl, hi) in enumerate(hlhis):
        print(f'\n[{TASK}, {task_name}] computing metrics for {hl}_{hi}...')
        output_file_for_head = output_dir.joinpath(f'{hl}_{hi}.csv')
        if output_file_for_head.is_file():
            return

        # load state dict for the attention head.
        state_dict_for_head = root_ckpt_path.joinpath(f'head_probing/{TASK}/{SUBTASK}{setting}/{shorthand}_{hl}/{shorthand}_{hl}_{hi}/')
        state_dict_for_head = list(state_dict_for_head.rglob("*.pt"))[0]
        state_dict_for_head = torch.load(state_dict_for_head)
        if TASK != 'NLI':
            state_dict_for_head = state_dict_for_head.state_dict()

        # then attach the probing layer
        model.attach_head_probe(hl, hi, n_classes=n_classes)

        # get the layer and check
        layer = model.network.get_attention_layer(hl)
        assert hasattr(layer, 'head_probe_dense_layer')

        # and load (put it on same device)
        if 'head_probe_dense_layer.weight' in state_dict_for_head.keys():
            weight = state_dict_for_head['head_probe_dense_layer.weight']
            bias = state_dict_for_head['head_probe_dense_layer.bias']
        else:
            weight = state_dict_for_head[f'bert.encoder.layer.{hl}.attention.self.head_probe_dense_layer.weight']
            bias = state_dict_for_head[f'bert.encoder.layer.{hl}.attention.self.head_probe_dense_layer.bias']

        layer.head_probe_dense_layer.weight = nn.Parameter(weight.to(device_id))
        layer.head_probe_dense_layer.bias = nn.Parameter(bias.to(device_id))

        # compute acc and save
        _, preds_for_layer, golds, ids = get_acc(model, test_data, metric_meta, device_id, head_probe=True)
        pd.Series(preds_for_layer).to_csv(output_file_for_head)

        # save labels and ids
        if not output_dir.joinpath(f'labels.csv').is_file():
            pd.Series(golds).to_csv(output_dir.joinpath(f'labels.csv'))

        if not output_dir.joinpath(f'ids.csv').is_file():
            pd.Series(ids).to_csv(output_dir.joinpath(f'ids.csv'))

        # detach
        model.detach_head_probe(hl)

def construct_model(checkpoint, task, task_def, return_base_model=False):
    task_defs = TaskDefs(task_def)
    assert task in task_defs._task_type_map
    assert task in task_defs._data_type_map
    assert task in task_defs._metric_meta_map

    prefix = task.split('_')[0]
    metric_meta = task_defs._metric_meta_map[task]

    # load model
    assert os.path.exists(checkpoint)

    if args.device_id > 0:
        state_dict = torch.load(checkpoint)
    else:
        state_dict = torch.load(checkpoint, map_location="cpu")

    config = state_dict['config']

    config['fp16'] = False
    config['answer_opt'] = 0
    config['adv_train'] = False

    task_def = task_defs.get_task_def(prefix)
    task_def_list = [task_def]
    config['task_def_list'] = task_def_list
    config["cuda"] = args.device_id >= 0
    config['device'] = args.device_id
    del state_dict['optimizer']

    model = MTDNNModel(config, device=config['device'])

    # scoring_list classification head doesn't matter because we're just taking
    # the head probe outputs.
    state_dict['state']['scoring_list.0.weight'] = model.network.state_dict()['scoring_list.0.weight']
    state_dict['state']['scoring_list.0.bias'] = model.network.state_dict()['scoring_list.0.bias']

    if not return_base_model:
        model.load_state_dict(state_dict)

    return model, metric_meta

def distribute_heads_to_gpus(available_devices):
    hlhis = []
    devices = []
    n_per_gpu = 144 // len(available_devices)
    gpu_group = []
    curr_device_idx = 0
    n_heads_cumul = 0

    for hl in range(12):
        for hi in range(12):
            n_heads_cumul += 1
            gpu_group.append((hl, hi))

            if ((n_heads_cumul % n_per_gpu == 0 and n_heads_cumul != 0) or n_heads_cumul == 144):
                hlhis.append(gpu_group)
                devices.append(available_devices[curr_device_idx])
                print(f'adding {len(gpu_group)} to device {available_devices[curr_device_idx]}')
                gpu_group = []
                curr_device_idx += 1

    return hlhis, devices

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available(),
                        help='whether to use GPU acceleration.')
    parser.add_argument('--device_id', type=int, default=0)
    parser.add_argument('--task', type=str)
    parser.add_argument('--subtask', type=str, default='head_probe')

    if False:
        models = [
            # 'ml_ltr'
            'base',
            'cross',
            'multi',
            # 'lang_spec/ml_ar',
            # 'lang_spec/ml_bg',
            # 'lang_spec/ml_de',
            # 'lang_spec/ml_el',
            # 'lang_spec/ml_es',
            # 'lang_spec/ml_fr',
            # 'lang_spec/ml_hi',
            # 'lang_spec/ml_ru',
            # 'lang_spec/ml_sw',
            # 'lang_spec/ml_th',
            # 'lang_spec/ml_tr',
            # 'lang_spec/ml_ur',
            # 'lang_spec/ml_vi',
            # 'lang_spec/ml_zh',
        ]

        if not SUBTASK == 'head_probe':
            single_lang_datasets = [
                'multi-ar',
                'multi-bg',
                'multi-de',
                'multi-el',
                'multi-es',
                'multi-fr',
                'multi-hi',
                'multi-ru',
                'multi-sw',
                'multi-th',
                'multi-tr',
                'multi-ur',
                'multi-vi',
                'multi-zh',
                'multi-en'
            ]

    else:
        models = [
            'cross',
            # 'multi',
            # 'base'
            ]
        assert SUBTASK == 'head_probe'

    root_ckpt_path = Path('checkpoint/')
    encoder_type = EncoderModelType.BERT

    for model_idx, model in enumerate(models):
        if TASK == 'NLI':
            if model in ['cross', 'multi', 'base']:
                task_def = f'experiments/{TASK}/{model}/task_def.yaml'
            else:
                lang = model.split("/")[1].split("_")[1]
                task_def = f'experiments/{TASK}/{f"multi-{lang}"}/task_def.yaml'

            if SUBTASK == 'cross_evaluate':
                assert model == 'multi' or model == 'multi-ltr'
                result_matrix = np.zeros((len(models), len(single_lang_datasets)))
                checkpoint = list(root_ckpt_path.joinpath(model).rglob("model_5_*.pt"))[0]

                model = construct_model(checkpoint, task_def, head_probe=False)
                accs = evaluate_model_against_multiple_datasets(
                    model,
                    single_lang_datasets,
                    task_def,
                    encoder_type,
                    args.batch_size,
                    args.max_seq_len,
                    device_id=args.device_id
                )
                for j, acc in enumerate(accs):
                    result_matrix[model_idx, j] = acc

            elif SUBTASK == 'head_probe':
                out_csv_file = f'score_outputs/NLI/head_probe/{model}/acc_per_head.csv'
                if not Path(out_csv_file).is_file():
                    data_path = f'experiments/NLI/cross/bert-base-multilingual-cased/cross_test.json'
                    result_matrix, result_df = evaluate_head_probe(
                        task_def,
                        encoder_type,
                        BATCH_SIZE,
                        MAX_SEQ_LEN,
                        device_id=args.device_id,
                        root_ckpt_path=root_ckpt_path,
                        setting=model,
                        data_path=data_path,
                        task_name=model,
                        return_base_model=model=='base'
                    )
                else:
                    result_df_all = pd.read_csv(out_csv_file, index_col=0)

        else:
            task_def = f'experiments/{TASK}/task_def.yaml'
            data_path = f'experiments/{TASK}/bert-base-multilingual-cased/{TASK.lower()}_test.json'
            out_pdf_file = f'score_outputs/{TASK}/head_probe/{model}/xnli_{model}-{TASK.lower()}'

            if not Path(out_pdf_file).is_file():
                available_devices = [0, 1, 2, 3]
                hlhis, devices = distribute_heads_to_gpus(available_devices)
                args = []

                for i, hlhi_set in enumerate(hlhis):
                    args_ = [
                        hlhi_set,
                        devices[i],
                        task_def,
                        encoder_type,
                        BATCH_SIZE,
                        MAX_SEQ_LEN,
                        root_ckpt_path,
                        model,
                        data_path,
                        TASK.lower(),
                        model=='base'
                    ]
                    args.append(tuple(args_))

                with torch.multiprocessing.get_context('spawn').Pool(len(available_devices)) as p:
                    p.starmap(evaluate_head_probe, args)

                # collect.
                output_dir = Path(f'score_outputs/{TASK}/head_probe/{model}')
                full_layer_results = []
                for hl in range(12):
                    for hi in range(12):
                        full_layer_results.append(pd.read_csv(output_dir.joinpath(f'{hl}_{hi}.csv'), index_col=0))

                full_ids = pd.read_csv(output_dir.joinpath(f'ids.csv'), index_col=0)
                full_golds = pd.read_csv(output_dir.joinpath(f'labels.csv'), index_col=0)

                full_results = pd.concat([full_ids, full_golds] + full_layer_results, axis=1)
                full_results.to_csv(output_dir.joinpath(f'results.csv'))