import json
import os
import shutil
from enum import Enum
from pathlib import Path
from typing import Tuple
import urllib.request
import sys
import time
import zipfile
from tqdm import tqdm
from distutils.dir_util import copy_tree

import pandas as pd
import numpy as np

import sdnist.strs as strs

DEFAULT_DATASET = 'diverse_communities_data_excerpts'

class TestDatasetName(Enum):
    NONE = 1
    GA_NC_SC_10Y_PUMS = 2
    NY_PA_10Y_PUMS = 3
    IL_OH_10Y_PUMS = 4
    taxi2016 = 5
    taxi2020 = 6
    ma2019 = 7
    tx2019 = 8
    national2019 = 9

dataset_name_state_map = {
    TestDatasetName.national2019.name: 'national',
    TestDatasetName.ma2019.name: 'massachusetts',
    TestDatasetName.tx2019.name: 'texas'
}

data_challenge_map = {
    TestDatasetName.NONE: None,
    TestDatasetName.GA_NC_SC_10Y_PUMS: strs.CENSUS,
    TestDatasetName.NY_PA_10Y_PUMS: strs.CENSUS,
    TestDatasetName.IL_OH_10Y_PUMS: strs.CENSUS,
    TestDatasetName.ma2019: strs.CENSUS,
    TestDatasetName.tx2019: strs.CENSUS,
    TestDatasetName.national2019: strs.CENSUS,
    TestDatasetName.taxi2016: strs.TAXI,
    TestDatasetName.taxi2020: strs.TAXI
}

def reporthook(count, block_size, total_size):
    global start_time
    if count == 0:
        start_time = time.time() - 1
        return
    duration = time.time() - start_time
    progress_size = int(count * block_size)
    speed = int(progress_size / (1024 * duration))
    progress = int(count * block_size * 100 / total_size)
    percent = min(progress, 100)  # keeps from exceeding 100% for small data
    sys.stdout.flush()
    sys.stdout.write("\r...%d%%, %d KB, %d KB/s, %d seconds elapsed" %
                     (percent, progress_size / 1024, speed, duration))

def error_opening_zip(zip_path):
    try:
        z = zipfile.ZipFile(zip_path)
    except zipfile.BadZipfile:
        return True

    if z.testzip() is not None:
        return True
    return False

def check_exists(root: Path, name: Path, download: bool, data_name: str = strs.DATA):
    root = root.expanduser()
    if not name.exists():
        print(f"{name} does not exist.")
        zip_path = Path(root.parent, 'data.zip')
        version = "2.3.0"

        version_v = f"v{version}"
        sdnist_version = DEFAULT_DATASET

        download_link = f"https://github.com/usnistgov/SDNist/releases/download/{version_v}/{sdnist_version}.zip"
        if zip_path.exists() and error_opening_zip(zip_path):
            os.remove(zip_path)

        if not zip_path.exists() and download:
            print(f"Downloading all SDNist datasets from: \n"
                  f"{download_link} ...")
            try:
                urllib.request.urlretrieve(
                    download_link,
                    zip_path.as_posix(),
                    reporthook
                )
                print(f'\n Success! Downloaded all datasets to "{root}" directory\n')
            except:
                if zip_path.exists():
                    shutil.rmtree(zip_path)
                raise RuntimeError(f"Unable to download {name}. Try: \n   "
                                   f"- re-running the command, \n   "
                                   f"- downloading manually from {download_link} "
                                   f"and unpack the zip. \n   "
                                   f"- or download the data as part of a release: https://github.com/usnistgov/SDNist/releases\n")

        if zip_path.exists():
            # extract zipfile
            extract_path = Path(root.parent, 'sdnist_data')
            with zipfile.ZipFile(zip_path, 'r') as f:
                for member in tqdm(f.infolist(), desc='Extracting '):
                    try:
                        f.extract(member, extract_path)
                    except zipfile.error as e:
                        raise e
            # delete zipfile
            os.remove(zip_path)
            print()
            copy_from_path = str(Path(extract_path, sdnist_version))
            copy_to_path = str(Path(root))
            print(f"Copying {copy_from_path} to {copy_to_path} ...")
            copy_tree(copy_from_path, copy_to_path)
            shutil.rmtree(extract_path)
        else:
            raise ValueError(f"{name} does not exist.")

def build_name(challenge: str,
               root: Path = Path("data"),
               public: bool = False,
               test: TestDatasetName = TestDatasetName.NONE,
               data_name: str = DEFAULT_DATASET):
    root = root.expanduser()
    directory = root

    if data_name == DEFAULT_DATASET:
        directory = root
    elif challenge == strs.CENSUS:
        directory = root / "census" / "dataset"
    elif challenge == strs.TAXI:
        directory = root / "taxi" / "dataset"

    if challenge == "census":
        if public:
            fname = "IL_OH_10Y_PUMS"

        elif test != TestDatasetName.NONE:
            if test in [TestDatasetName.taxi2020, TestDatasetName.taxi2016]:
                raise ValueError(f'Invalid challenge and test-dataset combination:'
                                 f'challenge: {challenge} | test-dataset: {test.name}. '
                                 f'Available test datasets for the census challenge: '
                                 f'{TestDatasetName.GA_NC_SC_10Y_PUMS.name}',
                                 f'{TestDatasetName.NY_PA_10Y_PUMS.name}')
            fname = test.name
        else:
            fname = TestDatasetName.NY_PA_10Y_PUMS.name

    elif challenge == "taxi":
        if public:
            fname = "taxi"
        elif test != TestDatasetName.NONE:
            if test in [TestDatasetName.GA_NC_SC_10Y_PUMS, TestDatasetName.NY_PA_10Y_PUMS]:
                raise ValueError(f'Invalid challenge and test-dataset combination: '
                                 f'challenge: {challenge} | test-dataset: {test.name}. '
                                 f'Available test datasets for the taxi challenge: '
                                 f'{TestDatasetName.taxi2016.name}',
                                 f'{TestDatasetName.taxi2020.name}')
            fname = test.name
        else:
            fname = TestDatasetName.taxi2020.name

    else:
        raise ValueError(f"Unrecognized challenge {challenge}")
    if fname in dataset_name_state_map.keys():
        fname = Path(dataset_name_state_map[fname], fname)
    return directory / fname

def load_parameters(challenge: str,
                    root: Path = Path("data"),
                    public: bool = True,
                    test: TestDatasetName = TestDatasetName.NONE,
                    download: bool = True,
                    data_name: str = DEFAULT_DATASET) -> dict:
    dataset_path = build_name(challenge=challenge, root=root,
                              public=public, test=test, data_name=data_name)
    dataset_parameters = dataset_path.with_suffix('.json')
    check_exists(root, dataset_parameters, download, data_name=data_name)

    with dataset_parameters.open("r") as handler:
        params = json.load(handler)
    return params

def load_config(challenge: str,
                root: Path = Path("data"),
                public: bool = True,
                test: TestDatasetName = TestDatasetName.NONE,
                download: bool = True) -> dict:
    dataset_path = build_name(challenge=challenge, root=root, public=public, test=test)
    dataset_config_path = dataset_path.with_suffix('.config.json')
    check_exists(root, dataset_config_path, download)

    with dataset_config_path.open("r") as handler:
        config = json.load(handler)

    if strs.K_MARGINAL in config:
        km_cnf = config[strs.K_MARGINAL]
        if strs.BINS in km_cnf and len(km_cnf[strs.BINS]):
            b_cnf = km_cnf[strs.BINS]
            for f, b in b_cnf.items():
                b_cnf[f] = eval(b)
    return config

def load_dataset(challenge: str,
                 root: Path = Path("data"),
                 public: bool = True,
                 test: TestDatasetName = TestDatasetName.NONE,
                 download: bool = True,
                 format_: str = "parquet",
                 data_name: str = DEFAULT_DATASET) -> Tuple[pd.DataFrame, dict]:
    """ Load one of the original SDNist datasets.

    :param challenge: str: base challenge. Must be `census` or `taxi`.
    :param root: Path: directory of the dataset.
    :param public: bool: whether to use the public dataset or the private dataset
        (see below).
    :param test: TestDatasetName: if not None, retrieve the given private dataset (see below).
    :param download: bool: download files if not present in the `root` directory.
    :param format_: str: preferred format when retrieving the files. Must be `parquet` or `csv`.
        Note that only `parquet` files are actually available for download.
    :param data_name: str: name of the dataset bundle from where to load data file
    :return: A tuple containing the requested dataset as a `pandas.DataFrame`, along with
        its corresponding schema, i.e a `dict` description of each feature of the dataset

    Regarding the public/private/test datasets:
    - during the challenge, the participants were given access to the "public"
    dataset (`public=True`, `test=False) to make experiments and inspect data.
    - their synthesizers were evaluated against two new datasets which they
    were not allowed to see before:
        1) The (`public=False`, `test=False`) corresponds to a dataset to privatize.
        The results were published on the public leaderboard. The participants
        were allowed to make several submission. In traditional machine learning scenarios,
        this would somewhat correspond to a 'validation' dataset.
        2) The (`public=False`, `test=True`) corresponds to another dataset to privatize.
        This time, the results were kepts secret until the end of the challenge.
        In traditional machine learning scenarios, this would correspond to a 'test' dataset.
        Your synthesizer should not be fine-tuned for the best score on this dataset.
    - using any information available in the 'public' dataset is allowed.
    - manually inspecting any of the two 'private' dataset is not allowed within the challenge,
        even though there is nothing preventing you from doing so.
    """
    if isinstance(root, str):
        root = Path(root)

    if public and test != TestDatasetName.NONE:
        raise ValueError("A public test dataset does not make sense.")

    name = build_name(challenge, root, public, test, data_name)
    name.parent.mkdir(exist_ok=True, parents=True)

    # Load schema
    params = load_parameters(challenge, root, public, test, download, data_name)
    config = dict()

    # TODO: remove .csv option

    # Load dataset
    if format_ == "parquet":
        dataset_name = name.with_suffix(".parquet")
        check_exists(root, dataset_name, download, data_name)
        dataset = pd.read_parquet(dataset_name)

    elif format_ == "csv":
        dataset_name = name.with_suffix(".csv")
        check_exists(root, dataset_name, download, data_name)

        columns = {name : params[strs.SCHEMA][name]["dtype"]
                   for name in params[strs.SCHEMA]}
        dataset = pd.read_csv(dataset_name, dtype=columns)

    else:
        raise ValueError(f"Unknown format {format_}")

    if data_name != DEFAULT_DATASET:
        config = load_config(challenge, root, public, test, download)
    params[strs.CONFIG] = config
    return dataset, params

if __name__ == "__main__":
    d = load_dataset("taxi")


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# coding: utf-8

"""
    Airflow API (Stable)

    Apache Airflow management API.  # noqa: E501

    The version of the OpenAPI document: 1.0.0
    Contact: dev@airflow.apache.org
    Generated by: https://openapi-generator.tech
"""

import pprint
import re  # noqa: F401

import six

from airflow_client.configuration import Configuration

class ListTaskInstanceForm(object):
    """NOTE: This class is auto generated by OpenAPI Generator.
    Ref: https://openapi-generator.tech

    Do not edit the class manually.
    """

    """
    Attributes:
      openapi_types (dict): The key is attribute name
                            and the value is attribute type.
      attribute_map (dict): The key is attribute name
                            and the value is json key in definition.
    """
    openapi_types = {
        'dag_ids': 'list[str]',
        'execution_date_gte': 'datetime',
        'execution_date_lte': 'datetime',
        'start_date_gte': 'datetime',
        'start_date_lte': 'datetime',
        'end_date_gte': 'datetime',
        'end_date_lte': 'datetime',
        'duration_gte': 'float',
        'duration_lte': 'float',
        'state': 'list[str]',
        'pool': 'list[str]',
        'queue': 'list[str]'
    }

    attribute_map = {
        'dag_ids': 'dag_ids',
        'execution_date_gte': 'execution_date_gte',
        'execution_date_lte': 'execution_date_lte',
        'start_date_gte': 'start_date_gte',
        'start_date_lte': 'start_date_lte',
        'end_date_gte': 'end_date_gte',
        'end_date_lte': 'end_date_lte',
        'duration_gte': 'duration_gte',
        'duration_lte': 'duration_lte',
        'state': 'State',
        'pool': 'Pool',
        'queue': 'Queue'
    }

    def __init__(self, dag_ids=None, execution_date_gte=None, execution_date_lte=None, start_date_gte=None, start_date_lte=None, end_date_gte=None, end_date_lte=None, duration_gte=None, duration_lte=None, state=None, pool=None, queue=None, local_vars_configuration=None):  # noqa: E501
        """ListTaskInstanceForm - a model defined in OpenAPI"""  # noqa: E501
        if local_vars_configuration is None:
            local_vars_configuration = Configuration()
        self.local_vars_configuration = local_vars_configuration

        self._dag_ids = None
        self._execution_date_gte = None
        self._execution_date_lte = None
        self._start_date_gte = None
        self._start_date_lte = None
        self._end_date_gte = None
        self._end_date_lte = None
        self._duration_gte = None
        self._duration_lte = None
        self._state = None
        self._pool = None
        self._queue = None
        self.discriminator = None

        if dag_ids is not None:
            self.dag_ids = dag_ids
        if execution_date_gte is not None:
            self.execution_date_gte = execution_date_gte
        if execution_date_lte is not None:
            self.execution_date_lte = execution_date_lte
        if start_date_gte is not None:
            self.start_date_gte = start_date_gte
        if start_date_lte is not None:
            self.start_date_lte = start_date_lte
        if end_date_gte is not None:
            self.end_date_gte = end_date_gte
        if end_date_lte is not None:
            self.end_date_lte = end_date_lte
        if duration_gte is not None:
            self.duration_gte = duration_gte
        if duration_lte is not None:
            self.duration_lte = duration_lte
        if state is not None:
            self.state = state
        if pool is not None:
            self.pool = pool
        if queue is not None:
            self.queue = queue

    @property
    def dag_ids(self):
        """Gets the dag_ids of this ListTaskInstanceForm.  # noqa: E501

        Return objects with specific DAG IDs. The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :return: The dag_ids of this ListTaskInstanceForm.  # noqa: E501
        :rtype: list[str]
        """
        return self._dag_ids

    @dag_ids.setter
    def dag_ids(self, dag_ids):
        """Sets the dag_ids of this ListTaskInstanceForm.

        Return objects with specific DAG IDs. The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :param dag_ids: The dag_ids of this ListTaskInstanceForm.  # noqa: E501
        :type: list[str]
        """

        self._dag_ids = dag_ids

    @property
    def execution_date_gte(self):
        """Gets the execution_date_gte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects greater or equal to the specified date. This can be combined with execution_date_lte parameter to receive only the selected period.   # noqa: E501

        :return: The execution_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._execution_date_gte

    @execution_date_gte.setter
    def execution_date_gte(self, execution_date_gte):
        """Sets the execution_date_gte of this ListTaskInstanceForm.

        Returns objects greater or equal to the specified date. This can be combined with execution_date_lte parameter to receive only the selected period.   # noqa: E501

        :param execution_date_gte: The execution_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._execution_date_gte = execution_date_gte

    @property
    def execution_date_lte(self):
        """Gets the execution_date_lte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects less than or equal to the specified date. This can be combined with execution_date_gte parameter to receive only the selected period.   # noqa: E501

        :return: The execution_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._execution_date_lte

    @execution_date_lte.setter
    def execution_date_lte(self, execution_date_lte):
        """Sets the execution_date_lte of this ListTaskInstanceForm.

        Returns objects less than or equal to the specified date. This can be combined with execution_date_gte parameter to receive only the selected period.   # noqa: E501

        :param execution_date_lte: The execution_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._execution_date_lte = execution_date_lte

    @property
    def start_date_gte(self):
        """Gets the start_date_gte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects greater or equal the specified date. This can be combined with startd_ate_lte parameter to receive only the selected period.   # noqa: E501

        :return: The start_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._start_date_gte

    @start_date_gte.setter
    def start_date_gte(self, start_date_gte):
        """Sets the start_date_gte of this ListTaskInstanceForm.

        Returns objects greater or equal the specified date. This can be combined with startd_ate_lte parameter to receive only the selected period.   # noqa: E501

        :param start_date_gte: The start_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._start_date_gte = start_date_gte

    @property
    def start_date_lte(self):
        """Gets the start_date_lte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects less or equal the specified date. This can be combined with start_date_gte parameter to receive only the selected period.   # noqa: E501

        :return: The start_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._start_date_lte

    @start_date_lte.setter
    def start_date_lte(self, start_date_lte):
        """Sets the start_date_lte of this ListTaskInstanceForm.

        Returns objects less or equal the specified date. This can be combined with start_date_gte parameter to receive only the selected period.   # noqa: E501

        :param start_date_lte: The start_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._start_date_lte = start_date_lte

    @property
    def end_date_gte(self):
        """Gets the end_date_gte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects greater or equal the specified date. This can be combined with start_date_lte parameter to receive only the selected period.   # noqa: E501

        :return: The end_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._end_date_gte

    @end_date_gte.setter
    def end_date_gte(self, end_date_gte):
        """Sets the end_date_gte of this ListTaskInstanceForm.

        Returns objects greater or equal the specified date. This can be combined with start_date_lte parameter to receive only the selected period.   # noqa: E501

        :param end_date_gte: The end_date_gte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._end_date_gte = end_date_gte

    @property
    def end_date_lte(self):
        """Gets the end_date_lte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects less than or equal to the specified date. This can be combined with start_date_gte parameter to receive only the selected period.   # noqa: E501

        :return: The end_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: datetime
        """
        return self._end_date_lte

    @end_date_lte.setter
    def end_date_lte(self, end_date_lte):
        """Sets the end_date_lte of this ListTaskInstanceForm.

        Returns objects less than or equal to the specified date. This can be combined with start_date_gte parameter to receive only the selected period.   # noqa: E501

        :param end_date_lte: The end_date_lte of this ListTaskInstanceForm.  # noqa: E501
        :type: datetime
        """

        self._end_date_lte = end_date_lte

    @property
    def duration_gte(self):
        """Gets the duration_gte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects greater than or equal to the specified values. This can be combined with duration_lte parameter to receive only the selected period.   # noqa: E501

        :return: The duration_gte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: float
        """
        return self._duration_gte

    @duration_gte.setter
    def duration_gte(self, duration_gte):
        """Sets the duration_gte of this ListTaskInstanceForm.

        Returns objects greater than or equal to the specified values. This can be combined with duration_lte parameter to receive only the selected period.   # noqa: E501

        :param duration_gte: The duration_gte of this ListTaskInstanceForm.  # noqa: E501
        :type: float
        """

        self._duration_gte = duration_gte

    @property
    def duration_lte(self):
        """Gets the duration_lte of this ListTaskInstanceForm.  # noqa: E501

        Returns objects less than or equal to the specified values. This can be combined with duration_gte parameter to receive only the selected range.   # noqa: E501

        :return: The duration_lte of this ListTaskInstanceForm.  # noqa: E501
        :rtype: float
        """
        return self._duration_lte

    @duration_lte.setter
    def duration_lte(self, duration_lte):
        """Sets the duration_lte of this ListTaskInstanceForm.

        Returns objects less than or equal to the specified values. This can be combined with duration_gte parameter to receive only the selected range.   # noqa: E501

        :param duration_lte: The duration_lte of this ListTaskInstanceForm.  # noqa: E501
        :type: float
        """

        self._duration_lte = duration_lte

    @property
    def state(self):
        """Gets the state of this ListTaskInstanceForm.  # noqa: E501

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :return: The state of this ListTaskInstanceForm.  # noqa: E501
        :rtype: list[str]
        """
        return self._state

    @state.setter
    def state(self, state):
        """Sets the state of this ListTaskInstanceForm.

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :param state: The state of this ListTaskInstanceForm.  # noqa: E501
        :type: list[str]
        """

        self._state = state

    @property
    def pool(self):
        """Gets the pool of this ListTaskInstanceForm.  # noqa: E501

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :return: The pool of this ListTaskInstanceForm.  # noqa: E501
        :rtype: list[str]
        """
        return self._pool

    @pool.setter
    def pool(self, pool):
        """Sets the pool of this ListTaskInstanceForm.

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :param pool: The pool of this ListTaskInstanceForm.  # noqa: E501
        :type: list[str]
        """

        self._pool = pool

    @property
    def queue(self):
        """Gets the queue of this ListTaskInstanceForm.  # noqa: E501

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :return: The queue of this ListTaskInstanceForm.  # noqa: E501
        :rtype: list[str]
        """
        return self._queue

    @queue.setter
    def queue(self, queue):
        """Sets the queue of this ListTaskInstanceForm.

        The value can be repeated to retrieve multiple matching values (OR condition).  # noqa: E501

        :param queue: The queue of this ListTaskInstanceForm.  # noqa: E501
        :type: list[str]
        """

        self._queue = queue

    def to_dict(self):
        """Returns the model properties as a dict"""
        result = {}

        for attr, _ in six.iteritems(self.openapi_types):
            value = getattr(self, attr)
            if isinstance(value, list):
                result[attr] = list(map(
                    lambda x: x.to_dict() if hasattr(x, "to_dict") else x,
                    value
                ))
            elif hasattr(value, "to_dict"):
                result[attr] = value.to_dict()
            elif isinstance(value, dict):
                result[attr] = dict(map(
                    lambda item: (item[0], item[1].to_dict())
                    if hasattr(item[1], "to_dict") else item,
                    value.items()
                ))
            else:
                result[attr] = value

        return result

    def to_str(self):
        """Returns the string representation of the model"""
        return pprint.pformat(self.to_dict())

    def __repr__(self):
        """For `print` and `pprint`"""
        return self.to_str()

    def __eq__(self, other):
        """Returns true if both objects are equal"""
        if not isinstance(other, ListTaskInstanceForm):
            return False

        return self.to_dict() == other.to_dict()

    def __ne__(self, other):
        """Returns true if both objects are not equal"""
        if not isinstance(other, ListTaskInstanceForm):
            return True

        return self.to_dict() != other.to_dict()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

import operator
import json
from collections import defaultdict, OrderedDict
from datetime import datetime

from django.http import Http404, HttpResponse
from django.shortcuts import \
    get_object_or_404, redirect, render, render_to_response
from django.contrib.auth.models import User
from django.contrib.auth.decorators import login_required, user_passes_test
from django.utils.decorators import method_decorator
from django.utils import timezone
from django.views.generic import ListView, DetailView
from django.db import transaction
from django.db.models import Q, Count

from .fdsn.fdsn_manager import \
    FdsnMotionManager, FdsnShakemapManager, \
    FdsnWaveformManager, FdsnDataselectManager, \
    OdcApiManager

from .models import \
    SearchEvent, SearchPeakMotions, \
    SearchCombined, SearchCustom
from .forms import \
    SearchEventsForm, SearchPeakMotionsForm, \
    SearchCombinedForm, SearchCustomForm

from .logger import RrsmLoggerMixin

class HomeListView(ListView):
    model = None
    context_object_name = 'events'
    template_name = 'events.html'

    def get_queryset(self):
        return None

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        motion_data, ws_url = FdsnMotionManager().get_event_list(days_back=31)
        context['breadcrumb'] = 'Home (events from last 31 days)'
        context['motion_data'] = motion_data
        context['ws_url'] = ws_url
        context['is_homepage'] = True
        return context

class AboutListView(ListView):
    model = None
    template_name = 'about.html'

    def get_queryset(self):
        return None

class RecentEventsListView(ListView):
    model = None
    context_object_name = 'events'
    template_name = 'events.html'

    def get_queryset(self):
        return None

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        motion_data, ws_url = FdsnMotionManager().get_event_list(
            days_back=int(self.kwargs.get('days'))
            )
        context['breadcrumb'] = 'Recent events'
        context['motion_data'] = motion_data
        context['ws_url'] = ws_url
        context['is_homepage'] = True
        return context

class EventDetailsListView(ListView, RrsmLoggerMixin):
    model = None
    context_object_name = None
    template_name = 'event_details.html'

    def __init__(self):
        RrsmLoggerMixin.__init__(self)

    def get_queryset(self):
        try:
            queryset = None
        except:
            raise
        return queryset

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        event_id = self.kwargs.get('event_public_id')
        fdsn_motion_man = FdsnMotionManager()
        fdsn_shakemap_man = FdsnShakemapManager()
        fdsn_waveform_man = FdsnWaveformManager()

        motion_data, ws_url = fdsn_motion_man.get_event_details(event_id)
        chart_pga, chart_pgv = self.get_sensor_channels_charts(motion_data)
        shakemap_url = fdsn_shakemap_man.get_shakemap_url(event_id)
        waveform_url = fdsn_waveform_man.get_waveform_url(event_id)

        context['motion_data'] = motion_data
        context['ws_url'] = ws_url
        context['chart_pga'] = chart_pga
        context['chart_pgv'] = chart_pgv
        context['shakemap_url'] = shakemap_url
        context['waveform_url'] = waveform_url
        return context

    def get_sensor_channels_charts(self, motion_data):
        try:
            if not motion_data:
                return None, None

            data_pga = defaultdict(dict)
            data_pgv = defaultdict(dict)
            chart_data_pga = {}
            chart_data_pgv = {}

            for s in motion_data.stations:
                for sc in s.sensor_channels:
                    data_pga[sc.channel_code][s.epicentral_distance] = \
                        sc.pga_value * 100
                    data_pgv[sc.channel_code][s.epicentral_distance] = \
                        sc.pgv_value * 100

            for key in data_pga:
                _tmp = sorted(data_pga[key].items())
                chart_data_pga[key] = _tmp

            for key in data_pgv:
                _tmp = sorted(data_pgv[key].items())
                chart_data_pgv[key] = _tmp

            chart_pga = {
                'exporting': {
                    'chartOptions': {
                        'plotOptions': {
                            'series': {
                            'dataLabels': {
                                'enabled': 'true'
                            }
                        }
                    }
                },
                'fallbackToExportServer': 'false'
            },
                'chart': {'type': 'scatter'},
                'title': {'text': 'PGA vs Epicentral distance'},
                'xAxis': {
                    'title': {
                        'text': 'Epicentral distance [km]'
                    },
                    'type': 'logarithmic',
                },
                'yAxis': {
                    'title': {
                        'text': 'PGA [cm/s^2]'
                    },
                    'type': 'logarithmic'
                },
                'series': []
            }

            chart_pgv = {
                'exporting': {
                    'chartOptions': {
                        'plotOptions': {
                            'series': {
                            'dataLabels': {
                                'enabled': 'true'
                            }
                        }
                    }
                },
                'fallbackToExportServer': 'false'
            },
                'chart': {'type': 'scatter'},
                'title': {'text': 'PGV vs Epicentral distance'},
                'xAxis': {
                    'title': {
                        'text': 'Epicentral distance [km]'
                    },
                    'type': 'logarithmic'
                },
                'yAxis': {
                    'title': {
                        'text': 'PGV [cm/s]'
                    },
                    'type': 'logarithmic'
                },
                'series': []
            }

            for t in chart_data_pga:
                chart_pga['series'].append({
                    'name': t,
                    'data': chart_data_pga[t]
                })

            for t in chart_data_pgv:
                chart_pgv['series'].append({
                    'name': t,
                    'data': chart_data_pgv[t]
                })

            dump_pga = json.dumps(chart_pga)
            dump_pgv = json.dumps(chart_pgv)
            return dump_pga, dump_pgv
        except:
            self.log_exception()
            return None, None

class StationStreamsListView(ListView, RrsmLoggerMixin):
    model = None
    context_object_name = None
    template_name = 'station_streams.html'

    def __init__(self):
        RrsmLoggerMixin.__init__(self)

    def get_queryset(self):
        try:
            queryset = None
        except:
            raise
        return queryset

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        event_id = self.kwargs.get('event_public_id')
        network_code = self.kwargs.get('network_code')
        station_code = self.kwargs.get('station_code')
        fdsn_motion_man = FdsnMotionManager()

        dataselect_url = FdsnDataselectManager(
            'NL', 'DBN', '2018-11-18T12:00:00'
        ).get_dataselect_url()

        motion_data, ws_url = fdsn_motion_man.get_event_details(
            event_id, network_code, station_code, True
        )
        chart_psa, chart_sd = self.get_spectra_charts(motion_data.stations[0])

        if motion_data:
            context['station_data'] = motion_data.stations[0]
        else:
            context['motion_data'] = None

        # waveform_pic = self.get_waveform_pic(motion_data)
        wafeform_picture = self.get_wafeform_picture(motion_data)

        context['chart_psa'] = chart_psa
        context['chart_sd'] = chart_sd
        # context['waveform_pic'] = waveform_pic
        context['ws_url'] = ws_url
        context['dataselect_url'] = dataselect_url
        context['wafeform_picture'] = wafeform_picture

        return context

    def get_spectra_charts(self, station_data):
        try:
            if not station_data:
                return None, None

            data_psa = defaultdict(dict)
            data_sd = defaultdict(dict)
            chart_data_psa = {}
            chart_data_sd = {}

            for sc in station_data.sensor_channels:
                for sa_psa in filter(
                    lambda x: x.type.lower() == 'psa', sc.spectral_amplitudes
                ):
                    data_psa[sc.channel_code][sa_psa.period] = sa_psa.amplitude
                for sa_sd in filter(
                    lambda x: x.type.lower() == 'sd', sc.spectral_amplitudes
                ):
                    data_sd[sc.channel_code][sa_sd.period] = sa_sd.amplitude

            for key in data_psa:
                _tmp = sorted(data_psa[key].items())
                chart_data_psa[key] = _tmp

            for key in data_sd:
                _tmp = sorted(data_sd[key].items())
                chart_data_sd[key] = _tmp

            chart_psa = {
                'exporting': {
                    'chartOptions': {
                        'plotOptions': {
                            'series': {
                            'dataLabels': {
                                'enabled': 'true'
                            }
                        },
                    }
                },
                'fallbackToExportServer': 'false'
            },
                'chart': {'type': 'scatter'},
                'plotOptions':{
                    'scatter':{
                        'lineWidth': '1'
                    }
                },
                'title': {'text': 'Pseudo-Spectral Acceleration'},
                'xAxis': {
                    'title': {
                        'text': 'Time [s]'
                    },
                    'type': 'logarithmic',
                },
                'yAxis': {
                    'title': {
                        'text': 'PSA [cm/s*s]'
                    },
                    'type': 'logarithmic'
                },
                'series': []
            }

            chart_sd = {
                'exporting': {
                    'chartOptions': {
                        'plotOptions': {
                            'series': {
                            'dataLabels': {
                                'enabled': 'true'
                            }
                        }
                    }
                },
                'fallbackToExportServer': 'false'
            },
                'chart': {'type': 'scatter'},
                'plotOptions':{
                    'scatter':{
                        'lineWidth': '1'
                    }
                },
                'title': {'text': 'Spectral Displacement Data'},
                'xAxis': {
                    'title': {
                        'text': 'Time [s]'
                    },
                    'type': 'logarithmic'
                },
                'yAxis': {
                    'title': {
                        'text': 'SD [cm]'
                    },
                    'type': 'logarithmic'
                },
                'series': []
            }

            for t in chart_data_psa:
                chart_psa['series'].append({
                    'name': t,
                    'data': chart_data_psa[t]
                })

            for t in chart_data_sd:
                chart_sd['series'].append({
                    'name': t,
                    'data': chart_data_sd[t]
                })

            dump_psa = json.dumps(chart_psa)
            dump_sd = json.dumps(chart_sd)
            return dump_psa, dump_sd
        except:
            self.log_exception()
            return None, None

    def get_wafeform_picture(self, station_data):
        try:
            data_wf = {}
            oam = OdcApiManager(station_data.stations[0])
            data = oam.get_waveform_data()

            if not data:
                return None

            for trace in data['payload'][0]['traces']:
                _tmp = []
                for e in trace['data']:
                    if not e[1]:
                        continue
                    # Mo need to convert the UNIX timestamp
                    # dt = datetime.utcfromtimestamp(int(e[0])/1000)
                    _tmp.append([e[0], e[1]])
                data_wf[trace['name']] = _tmp

            chart = {
                'exporting': {
                    'chartOptions': {
                        'plotOptions': {
                            'series': {
                            'dataLabels': {
                                'enabled': 'true'
                            }
                        }
                    }
                },
                'fallbackToExportServer': 'false'
            },
                'chart': {'type': 'spline'},
                'title': {'text': 'Waveform'},
                'plotOptions': {
                    'series': {
                        'lineWidth': 1
                    }
                },
                'xAxis': {
                    'type': 'datetime',
                    'labels': {
                        'format': '{value:%H:%M:%S}'
                    }
                },
                'series': []
            }

            for t in data_wf:
                chart['series'].append({
                    'name': t,
                    'data': data_wf[t]
                })

            dump = json.dumps(chart)
            return dump
        except:
            raise

def search_events(request):
    if request.method == 'POST':
        form = SearchEventsForm(request.POST)
        if form.is_valid():
            data, ws_url = FdsnMotionManager().get_event_list(
                event_id=form.cleaned_data['event_id'],
                date_start=form.cleaned_data['date_start'],
                date_end=form.cleaned_data['date_end'],
                magnitude_min=form.cleaned_data['magnitude_min'],
                network_code=form.cleaned_data['network_code'],
                station_code=form.cleaned_data['station_code'],
                event_lat_min=form.cleaned_data['event_lat_min'],
                event_lat_max=form.cleaned_data['event_lat_max'],
                event_lon_min=form.cleaned_data['event_lon_min'],
                event_lon_max=form.cleaned_data['event_lon_max']
            )

            return render(
                request,
                'events.html',
                {
                    'motion_data': data,
                    'breadcrumb': 'Search result (events)',
                    'ws_url': ws_url,
                    'is_homepage': True
                }
            )
        else:
            return render(
                request,
                'search.html',
                {
                    'form': form,
                    'show_info': True
                }
            )
    else:
        form = SearchEventsForm(instance=SearchEvent())
        return render(
            request,
            'search.html',
            {
                'form': form,
                'show_info': True
            }
        )

def search_peak_motions(request):
    if request.method == 'POST':
        form = SearchPeakMotionsForm(request.POST)
        if form.is_valid():
            data, ws_url = FdsnMotionManager().get_event_list(
                pga_min=form.cleaned_data['pga_min'],
                pga_max=form.cleaned_data['pga_max'],
                pgv_min=form.cleaned_data['pgv_min'],
                pgv_max=form.cleaned_data['pgv_max']
            )

            return render(
                request,
                'events.html',
                {
                    'motion_data': data,
                    'breadcrumb': 'Search result (peak motions)',
                    'ws_url': ws_url,
                    'is_homepage': True
                }
            )
        else:
            return render(
                request,
                'search.html',
                {
                    'form': form
                }
            )
    else:
        form = SearchPeakMotionsForm(instance=SearchPeakMotions())
        return render(
            request,
            'search.html',
            {
                'form': form
            }
        )

def search_combined(request):
    if request.method == 'POST':
        form = SearchCombinedForm(request.POST)
        if form.is_valid():
            data, ws_url = FdsnMotionManager().get_event_list(
                magnitude_min=form.cleaned_data['magnitude_min'],
                pga_min=form.cleaned_data['pga_min'],
                pga_max=form.cleaned_data['pga_max'],
                pgv_min=form.cleaned_data['pgv_min'],
                pgv_max=form.cleaned_data['pgv_max'],
                stat_lat_min=form.cleaned_data['stat_lat_min'],
                stat_lat_max=form.cleaned_data['stat_lat_max'],
                stat_lon_min=form.cleaned_data['stat_lon_min'],
                stat_lon_max=form.cleaned_data['stat_lon_max']
            )

            return render(
                request,
                'events.html',
                {
                    'motion_data': data,
                    'breadcrumb': 'Search result (combined)',
                    'ws_url': ws_url,
                    'is_homepage': True
                }
            )
        else:
            return render(
                request,
                'search.html',
                {
                    'form': form,
                    'show_info': True
                }
            )
    else:
        form = SearchCombinedForm(instance=SearchCombined())
        return render(
            request,
            'search.html',
            {
                'form': form,
                'show_info': True
            }
        )

def search_custom(request):
    if request.method == 'POST':
        form = SearchCustomForm(request.POST)
        if form.is_valid():
            data, ws_url = FdsnMotionManager().get_event_list(
                event_id=form.cleaned_data['event_id'],
                date_start=form.cleaned_data['date_start'],
                date_end=form.cleaned_data['date_end'],
                magnitude_min=form.cleaned_data['magnitude_min'],
                network_code=form.cleaned_data['network_code'],
                station_code=form.cleaned_data['station_code'],
                pga_min=form.cleaned_data['pga_min'],
                pga_max=form.cleaned_data['pga_max'],
                pgv_min=form.cleaned_data['pgv_min'],
                pgv_max=form.cleaned_data['pgv_max'],
                stat_lat_min=form.cleaned_data['stat_lat_min'],
                stat_lat_max=form.cleaned_data['stat_lat_max'],
                stat_lon_min=form.cleaned_data['stat_lon_min'],
                stat_lon_max=form.cleaned_data['stat_lon_max'],
                event_lat_min=form.cleaned_data['event_lat_min'],
                event_lat_max=form.cleaned_data['event_lat_max'],
                event_lon_min=form.cleaned_data['event_lon_min'],
                event_lon_max=form.cleaned_data['event_lon_max']
            )

            return render(
                request,
                'events.html',
                {
                    'motion_data': data,
                    'breadcrumb': 'Search result (custom)',
                    'ws_url': ws_url,
                    'is_homepage': True
                }
            )
        else:
            return render(
                request,
                'search_custom.html',
                {
                    'form': form
                }
            )
    else:
        form = SearchCustomForm(instance=SearchCustom())
        return render(
            request,
            'search_custom.html',
            {
                'form': form
            }
        )

def download_waveforms(request):
    if request.method == 'POST':
        urls = request.POST.get('urls')
        return HttpResponse(
            urls,
            content_type="text/plain"
        )
    else:
        return None

def custom_404(request, exception=None):
    '''HTTP 404 custom handler
    '''
    return render_to_response('404.html', {'exception': exception})

def custom_500(request, exception=None):
    '''HTTP 500 custom handler
    '''
    return render_to_response('500.html', {'exception': exception})


from PIL import Image
import math
im = Image.open("sample_image.jpg")
circumRadius = 17
cr = circumRadius
oR = round(cr / math.cos(math.radians(30)))

f = open("sample_image.svg", "w+")
f.truncate()
f.write('<?xml version="1.0" standalone="yes"?>' + '\n')
f.write('<parent xmlns="http://example.org" xmlns:svg="http://www.w3.org/2000/svg">' + '\n')
f.write('<svg:svg width="' + str(im.size[0]) + '" height="' + str(im.size[1]) + '" version="1.1">' + '\n')

def paintHexShape(cX, cY, cr, colour):
    points = [
        str(round(cX - cr / 2)) + ',' + str(cY - cr),
        str(round(cX + cr / 2)) + ',' + str(cY - cr),
        str(cX + oR) + ',' + str(cY),
        str(round(cX + cr / 2)) + ',' + str(cY + cr),
        str(round(cX - cr / 2)) + ',' + str(cY + cr),
        str(cX - oR) + ',' + str(cY)
    ]
    s = ' '
    pointsString = s.join(points)
    f.write('<svg:polygon fill="'  + colour + '" points="' + pointsString + '" />' + '\n')

def getAvgHexColour(x, y, cr):
    widthsTopHalf = [j * cr for j in range(1, cr + 1)]
    for i in range(cr):
        widthsTopHalf[i] += i
    widthsBotHalf = widthsTopHalf[::-1]
    allWidths = widthsTopHalf + widthsBotHalf
    colours = list()
    for i in range(cr * 2):
        tX = x - allWidths[i] / 2
        tY = y - cr + i
        ti = im.crop((tX, tY, tX + allWidths[i], tY + 1)).resize((1, 1))
        colours.append(ti.getpixel((0,0)))
    widthWeight = [i / allWidths[cr] for i in allWidths]
    returnColour = [0, 0, 0]
    for j in range(3):
        returnColour[j] = round(sum([colours[i][j] * widthWeight[i] for i in range(len(colours))]) / (cr * 2))
    tuple(returnColour)
    return '#{:02x}{:02x}{:02x}'.format(*returnColour)

def hexShapes(im, cr):
    startPointsX = [i for i in range(0, im.size[0] +  cr, round(oR + (cr / 2)) - 1)]
    startPointsY = [j for j in range(0, im.size[1] +  cr, 2 * cr - 1)]

    oddRow = True
    for i in startPointsX:
        for j in startPointsY:
            colour = getAvgHexColour(i, j, cr)
            if(oddRow):
                paintHexShape(i, j, cr, colour)
            else:
                paintHexShape(i, j - cr, cr, colour)
            print("Hexagon at " + str(i) + ',' + str(j) + " created")
        oddRow = not oddRow

hexShapes(im, cr)

f.write('</svg:svg>' + '\n')
f.write('</parent>')


#!/usr/bin/python3

import json
import cgi
import traceback
from DBController import DBController
import re

def get_resp():

   def PrintRow(name, data, col, mode, REGIONS, TIMEZONES, OPERATORS):
      res = "{}	".format(name)

      if mode == 'TimeZoneGroup':
         for TimeZone in TIMEZONES:
            res += """{}	""".format(data[TimeZone][col])

      if mode == 'RegionGroup':
         for Region in REGIONS:
            res += """{}	""".format(data[Region][col])

      if mode == 'OperatorGroup':
         for operator in OPERATORS:
            res += """{}	""".format(data[operator][col])

      res += """{}

""".format(str(data["res"][col]))

      return res

   def PrintRowDiv(name, data, col, base, mode, REGIONS, TIMEZONES, OPERATORS):
      res = "{}	".format(name)

      if mode == 'TimeZoneGroup':
         for TimeZone in TIMEZONES:
            if data[TimeZone][base] != 0:
               res += """{}	""".format(str(data[TimeZone][col]) +
                      " ({:.1f}%)".format(data[TimeZone][col]/data[TimeZone][base] * 100))
            else:
               res += """0	"""

      if mode == 'RegionGroup':
         for Region in REGIONS:
            if data[Region][base] != 0:
               res += """{}	""".format(str(data[Region][col]) +
                      " ({:.1f}%)".format(data[Region][col]/data[Region][base] * 100))
            else:
               res += """0	"""

      if mode == 'OperatorGroup':
         for operator in OPERATORS:
            if data[operator][base] != 0:
               res += """{}	""".format(str(data[operator][col]) +
                      " ({:.1f}%)".format(data[operator][col]/data[operator][base] * 100))
            else:
               res += """0	"""

      if data["res"][base] != 0:
         res += """{}

""".format(str(data["res"][col]) + " ({:.1f}%)".format(data["res"][col]/data["res"][base] * 100))
      else:
         res += """0	"""

      return res

   def PrintDetail(data, info, TEMPLATE, mode, REGIONS, TIMEZONES, OPERATORS, is_bold):

      style = ""

      if is_bold:
         style = "font-weight: bold;font-size: larger;"

      res = """

               {}	""".format(style,str(len(TEMPLATE)),data)

      res += PrintRow("Всего звонков", info[data], "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Отправлено", info[data], "SEND", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Сделок", info[data], "BID", "SEND", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Повторная обработка", info[data], "RECALL", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Плохие данные", info[data], "BADDATA", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Отказ", info[data], "NEGATIV", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("ОДП", info[data], "ODP", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Уникальных контактов", info[data], "UNIQ_COUNT", "COUNT", mode, REGIONS, TIMEZONES, OPERATORS)
      res += '
'.format(style) + PrintRowDiv("Отправлено уникальных", info[data], "UNIQ_SEND", "UNIQ_COUNT", mode, REGIONS, TIMEZONES, OPERATORS)

      return res

   form = cgi.FieldStorage()
   mode = form.getvalue('mode')
   start_date = form.getvalue('start_date')
   end_date = form.getvalue('end_date')
   detail = form.getvalue('detail')
   controller = DBController()
   cur = controller.GetCur()

   main_sql = """SELECT """

   if detail == "Month":
      main_sql += 'DATE_FORMAT(DateTime, "%m.%Y") as Detail,'
   elif detail == "Hours":
      main_sql += 'DATE_FORMAT(DateTime, "%d.%m.%Y %H") as Detail,'
   else:
      main_sql += 'DATE_FORMAT(DateTime, "%d.%m.%Y") as Detail,'

   main_sql += """ Result, IF(PrevCall_ IS NULL, false, true) as uniq, IF(Bid.State = 'Счет открыт', true, false) as BID, count(*) as CNT"""

   if mode == 'OperatorGroup':
      main_sql += ", CONCAT(Employee.Name, ' ', Employee.Surname) as grp"

   if mode == 'TimeZoneGroup':
      main_sql += ', TimeZone as grp'

   if mode == 'RegionGroup':
      main_sql += ', TimeZone.Region as grp'

   main_sql += """
FROM Call_ LEFT JOIN Bid ON Call_.ID = Bid.Call_ """

   if mode == 'OperatorGroup':
      main_sql += """INNER JOIN Operator ON Call_.Operator = Operator.ID
                     INNER JOIN Employee ON Operator.Employee = Employee.ID"""

   if mode == 'TimeZoneGroup' or mode == 'RegionGroup':
      main_sql += """INNER JOIN Call_Organization ON Call_Organization.Call_ = Call_.ID
                     INNER JOIN Organization ON Call_Organization.Organization = Organization.ID"""

   if mode == 'RegionGroup':
      main_sql += ' LEFT JOIN TimeZone ON SUBSTR(Organization.Address,1,3) >= TimeZone.StartValue AND ' \
                  'SUBSTR(Organization.Address,1,3) <= TimeZone.EndValue'

   main_sql += """
WHERE DateTime >= '{}' AND DateTime <= '{}' AND Result is NOT NULL AND Result <> 'ОДПРоботТФ'
GROUP BY """

   if detail == "Month":
      main_sql += 'MONTH( DateTime ), YEAR(DateTime)'
   elif detail == "Hours":
      main_sql += 'Date(DateTime), HOUR(DateTime)'
   else:
      main_sql += 'Date(DateTime)'

   main_sql += """, Result, IF(PrevCall_ IS NULL, false, true), IF(Bid.State = 'Счет открыт', true, false)"""

   if mode == 'OperatorGroup':
      main_sql += ", Employee.Name, Employee.Surname"

   if mode == 'TimeZoneGroup':
      main_sql += ', TimeZone'

   if mode == 'RegionGroup':
      main_sql += ', TimeZone.Region'

   cur.execute(main_sql.format(start_date,end_date))

   res = cur.fetchall()

   OPERATORS = []
   if mode == 'OperatorGroup':
      cur.execute("""SELECT Employee.Surname, Employee.`Name`
                     FROM Operator
                     INNER JOIN Employee ON Operator.Employee = Employee.ID
                     WHERE Dismissed = FALSE""")

      op_list = cur.fetchall()

      for data in res:
         if not data['grp'] in OPERATORS:
            OPERATORS.append( data['grp'] )

   info = {}
   row_res = None

   REGIONS = ["01","02","03","04","05","06","07","08","09","10","11","12","13","14","15","16",
              "17","18","19","21","22","23","24","25","26","27","28","29","30","31","32","33",
              "34","35","36","37","38","39","40","41","42","43","44","45","46","47","48","49",
              "50","51","52","53","54","55","56","57","58","59","60","61","62","63","64","65",
              "66","67","68","69","70","71","72","73","74","75","76","77","78","79","83","86",
              "87","89","91","92","95",None]

   TIMEZONES = [9,8,7,6,5,4,3,2,1,0,-1,None]

   TEMPLATE = {"COUNT":0, "SEND":0, "RECALL":0, "BADDATA":0, "NEGATIV":0, "ODP":0, "UNIQ_COUNT":0, "UNIQ_SEND":0, "BID":0}

   if mode == "NonGroup":
      row_res = TEMPLATE.copy()

   if mode == 'TimeZoneGroup':
      row_res = {}

      for TimeZone in TIMEZONES:
         row_res[TimeZone]= TEMPLATE.copy()

   if mode == 'RegionGroup':
      row_res = {}

      for Region in REGIONS:
         row_res[Region]= TEMPLATE.copy()

   if mode == 'OperatorGroup':
      row_res = {}

      for operator in OPERATORS:
         row_res[operator]= TEMPLATE.copy()

   glob_info = TEMPLATE.copy()

   for data in res:
      cur_row = None

      if not data["Detail"] in info:
         info[data["Detail"]] = {}
         cur_row = info[data["Detail"]]
         cur_row["res"] = TEMPLATE.copy()

         if mode == 'TimeZoneGroup':
            for TimeZone in TIMEZONES:
               cur_row[TimeZone]= TEMPLATE.copy()

         if mode == 'RegionGroup':
            for Region in REGIONS:
               cur_row[Region]= TEMPLATE.copy()

         if mode == 'OperatorGroup':
            for operator in OPERATORS:
               cur_row[operator]= TEMPLATE.copy()
      else:
         cur_row = info[data["Detail"]]

      cur_row["res"]["COUNT"] += data["CNT"]
      glob_info["COUNT"] += data["CNT"]

      if "grp" in data:
         row_res[data["grp"]]["COUNT"] += data["CNT"]
         cur_row[data["grp"]]["COUNT"] += data["CNT"]

      if data["Result"] == 'Отправлено':
         cur_row["res"]["SEND"] += data["CNT"]
         glob_info["SEND"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["SEND"] += data["CNT"]
            cur_row[data["grp"]]["SEND"] += data["CNT"]

      if data["Result"] in ('Недозвон', 'Повторный звонок', 'Отказ. Пока не актуально', 'Дальний регион'):
         cur_row["res"]["RECALL"] += data["CNT"]
         glob_info["RECALL"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["RECALL"] += data["CNT"]
            cur_row[data["grp"]]["RECALL"] += data["CNT"]

      if data["Result"] in ('Неактивен', 'Неправильный номер', 'Дубль', 'Отказ. Крым'):
         cur_row["res"]["BADDATA"] += data["CNT"]
         glob_info["BADDATA"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["BADDATA"] += data["CNT"]
            cur_row[data["grp"]]["BADDATA"] += data["CNT"]

      if data["Result"] in ('ОДП'):
         cur_row["res"]["ODP"] += data["CNT"]
         glob_info["ODP"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["ODP"] += data["CNT"]
            cur_row[data["grp"]]["ODP"] += data["CNT"]

      if data["Result"] in ('Отказ. Сбер', 'Отказ. Точка', 'Отказ. ТФ', 'Отказ. Др.банк', 'Отказ. Открыл', 'Отказ. Звонили', 'Отказ. Нет потребности', 'Отказ. Услуга не нужна', 'Отказ. Мск'):
         cur_row["res"]["NEGATIV"] += data["CNT"]
         glob_info["NEGATIV"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["NEGATIV"] += data["CNT"]
            cur_row[data["grp"]]["NEGATIV"] += data["CNT"]

      if data["uniq"]:
         cur_row["res"]["UNIQ_COUNT"] += data["CNT"]
         glob_info["UNIQ_COUNT"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["UNIQ_COUNT"] += data["CNT"]
            cur_row[data["grp"]]["UNIQ_COUNT"] += data["CNT"]

      if data["uniq"] and data["Result"] == 'Отправлено':
         cur_row["res"]["UNIQ_SEND"] += data["CNT"]
         glob_info["UNIQ_SEND"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["UNIQ_SEND"] += data["CNT"]
            cur_row[data["grp"]]["UNIQ_SEND"] += data["CNT"]

      if data["BID"]:
         cur_row["res"]["BID"] += data["CNT"]
         glob_info["BID"] += data["CNT"]

         if "grp" in data:
            row_res[data["grp"]]["BID"] += data["CNT"]
            cur_row[data["grp"]]["BID"] += data["CNT"]

   res = {}
   res["statContent"] = """

Дата звонка
	"""

   if mode == 'TimeZoneGroup':
      for TimeZone in TIMEZONES:
         res["statContent"] += """{}	""".format(TimeZone)

   if mode == 'RegionGroup':
      for Region in REGIONS:
         res["statContent"] += """{}	""".format(Region)

   if mode == 'OperatorGroup':
      for operator in OPERATORS:
         res["statContent"] += """{}	""".format(operator)

   res["statContent"] += """Итого

"""

   for data in info:
      res["statContent"] += PrintDetail(data, info, TEMPLATE, mode, REGIONS, TIMEZONES, OPERATORS, False)

   row_res['res'] = glob_info
   info['Итого'] = row_res
   res["statContent"] += PrintDetail('Итого', info, TEMPLATE, mode, REGIONS, TIMEZONES, OPERATORS, True)

   res["statContent"] += "
"
   print("Content-type: application/json")
   print()
   print(json.JSONEncoder().encode(res))

try:
   get_resp()
except Exception as e:
   res = {}
   res["statContent"] = "\nОшибка:\n" + str(e) + "\n" + traceback.format_exc()
   print("Content-type: application/json")
   print()
   print(json.JSONEncoder().encode(res))


import geopandas as gpd
from loguru import logger

from dublin_building_stock.spatial_operations import (
    get_geometries_within,
)

def create_dublin_small_area_boundaries(data_dir, dublin_boundary):

    filenames = [
        "Census2011_Small_Areas_generalised20m",
        "Small_Areas_Ungeneralised_-_OSi_National_Statistical_Boundaries_-_2015-shp",
    ]
    filepaths = [data_dir / filename for filename in filenames]
    for filepath in filepaths:
        use_columns = ["SMALL_AREA", "EDNAME", "geometry"]
        ireland_small_area_boundaries = gpd.read_file(filepath)[use_columns]
        dublin_small_area_boundaries = get_geometries_within(
            ireland_small_area_boundaries.to_crs(epsg=2157),
            dublin_boundary.to_crs(epsg=2157),
        )
        output_filepath = filepath.parent / ("Dublin_" + filepath.stem)
        dublin_small_area_boundaries.to_file(output_filepath)


import lasagne
from lasagne.layers import DenseLayer, NonlinearityLayer

from ..utils.format import check_list
from ..utils.layers import clip_grads, add, mul
from ..memory.gate import GateLayer

# Vanilla RNN cell

def RNNCell(prev_state,
            input_or_inputs=tuple(),
            nonlinearity=lasagne.nonlinearities.sigmoid,
            num_units=None,
            name="YetAnotherRNNLayer",
            grad_clipping=5.):
    """

    Implements a one-step recurrent neural network (RNN) with arbitrary number of units.

    parameters:
        prev_state: input that denotes previous state (shape must be (None, n_units) )
        input_or_inputs: a single layer or a list/tuple of layers that go as inputs
        nonlinearity: which nonlinearity to use
        num_units: how many recurrent cells to use. None means "as in prev_state"
        grad_clipping: maximum gradient absolute value. 0 or None means "no clipping"

    for developers:
        Works by stacking DenseLayers with ElemwiseSumLayer.
        is a function mock, not actual class.

    """

    assert len(prev_state.output_shape) == 2
    # if needed, infer num_units
    if num_units is None:
        num_units = prev_state.output_shape[1]
    # else check it
    assert num_units == prev_state.output_shape[1]

    inputs = check_list(input_or_inputs)

    if grad_clipping:
        prev_state = clip_grads(prev_state, grad_clipping)
        inputs = [clip_grads(lyr, grad_clipping) for lyr in inputs]

    # from prev state to current state (with bias)
    hid_to_hid = DenseLayer(prev_state,
                            num_units=num_units,
                            nonlinearity=None,
                            name=name + ".hid_to_hid")

    # from inputs to current state (without bias)
    inputs_to_hid = [DenseLayer(input_layer,
                                num_units=num_units,
                                nonlinearity=None,
                                b=None,  # Cruicial! This disables additional bias layers
                                name=name + ".input%i_to_hid" % (i))
                     for i, input_layer in enumerate(inputs)]

    # stack them
    elwise_sum = add(*([hid_to_hid] + inputs_to_hid), name=name + ".sum")

    # finally, apply nonlinearity

    new_hid = NonlinearityLayer(elwise_sum,
                                nonlinearity,
                                name=name + ".new_state")

    return new_hid

def GRUCell(prev_state,
            input_or_inputs=tuple(),
            num_units=None,
            forgetgate_nonlinearity=lasagne.nonlinearities.sigmoid,
            updategate_nonlinearity=lasagne.nonlinearities.sigmoid,
            hidden_update_nonlinearity=lasagne.nonlinearities.tanh,
            name="YetAnotherGRULayer",
            grad_clipping=5.
            ):
    """

    Implements a one-step gated recurrent unit (GRU) with arbitrary number of units.

    parameters:
        prev_state: input that denotes previous state (shape must be (None, n_units) )
        input_or_inputs: a single layer or a list/tuple of layers that go as inputs
        *_nonlinearity: which nonlinearity to use for a particular gate
        num_units: how many recurrent cells to use. None means "as in prev_state"
        grad_clipping: maximum gradient absolute value. 0 or None means "no clipping"

    for developers:
        Works by stacking other lasagne layers;
        is a function mock, not actual class.

    """

    assert len(prev_state.output_shape) == 2
    # if required, infer num_units
    if num_units is None:
        num_units = prev_state.output_shape[1]
    # else check it
    assert num_units == prev_state.output_shape[1]

    inputs = check_list(input_or_inputs)

    # hidden to gates
    hid_to_gates = GateLayer(prev_state, [num_units] * 3,
                             gate_nonlinearities=None,
                             bias_init=None,
                             name=name + ".hidden_to_gates_stacked")

    hid_forget, hid_update, hidden_update_hid = hid_to_gates.values()

    # clip grads #1
    if grad_clipping:
        inputs = [clip_grads(lyr, grad_clipping) for lyr in inputs]
        hid_forget, hid_update, hidden_update_hid = [clip_grads(lyr, grad_clipping) for lyr in
                                                     [hid_forget, hid_update, hidden_update_hid]]

    # input to gates
    inp_to_gates = GateLayer(inputs, [num_units] * 3,
                             gate_nonlinearities=None,
                             name=name + ".input_to_gates_stacked")
    inp_forget, inp_update, hidden_update_in = inp_to_gates.values()

    # compute forget and update gates
    forgetgate = NonlinearityLayer(
        add(inp_forget, hid_forget),
        forgetgate_nonlinearity,
        name="forgetgate"
    )
    updategate = NonlinearityLayer(
        add(inp_update, hid_update),
        updategate_nonlinearity,
        name="updategate"
    )

    inv_updategate = NonlinearityLayer(updategate,
                               lambda x: 1 - x,
                               name="1 - updategate")

    # compute hidden update
    hidden_update = add(
        hidden_update_in,
        mul(forgetgate, hidden_update_hid),
        name="hid_update"
    )

    # clip grads #2
    if grad_clipping:
        hidden_update = clip_grads(hidden_update,
                                   grad_clipping)

    hidden_update = NonlinearityLayer(hidden_update,
                                      hidden_update_nonlinearity)

    # compute new hidden values
    new_hid = add(
        mul(inv_updategate, prev_state),
        mul(updategate, hidden_update),
        name=name
    )

    return new_hid


import numpy as np
import matplotlib.pyplot as plt
from data_processing_func import leave_enter_shelter
import plotting_functions

def get_flights(session, startf):
    """
    A flight is defined as an event in which the mouse leaves the threat zone and runs to the shelter zone. There is no
    threshold on the speed but the movement must be without interruptions
    """
    # Prepare data, get frames at which mouse enters and leaves rois
    transitions = Transitions()

    threat_inroi_idxs = session.inroi.inthreat_indexes.copy()
    threat_in = [t for s, t in zip(threat_inroi_idxs, threat_inroi_idxs[1:]) if t - s > 90]
    threat_in.append(threat_inroi_idxs[0])
    threat_in = sorted(threat_in)
    threat_out = [threat_inroi_idxs[threat_inroi_idxs.index(x)-1]+1 for idx,x in enumerate(threat_in)][1::]

    shelter_inroir_idxs = session.inroi.inshelter_indexes.copy()
    shelter_in = [t for s, t in zip(shelter_inroir_idxs, shelter_inroir_idxs[1:]) if t - s > 90]
    shelter_in.append(shelter_inroir_idxs[0])
    shelter_in = sorted(shelter_in)
    shelter_out = [shelter_inroir_idxs[shelter_inroir_idxs.index(x)-1]+1 for idx,x in enumerate(shelter_in)][1::]

    # Get transition times
    trans_s_to_t = []
    for s in shelter_out:
        if s > startf:
            # Get timepoint at which mouse enters threat roi
            if not [x for x in session.tracking.x[s:-1] if x > session.tracking.threat[0]]:
                continue
            frames_to_threat = [idx for idx,x in enumerate(session.tracking.x[s:-1]) if x > session.tracking.threat[0]][0]

            # check if before that there is another shelter roi exit
            if not [x for x in shelter_out if s<x<s+frames_to_threat]:
                # Get the time at which the mouse last entered the shelter roi
                last_entrance = [x for x in shelter_in if x < s][-1]
                # store the results
                trans_s_to_t.append((s, s + frames_to_threat))

    # now repeat from threat roi to shelter roi transitions
    trans_t_to_s = []
    for t in threat_out:
        if t > startf:
            if not [x for x in session.tracking.x[t:-1] if x > session.tracking.threat[0]+session.tracking.shelter[2]]:
                continue

            if [idx for idx,x in enumerate(session.tracking.x[t:-1])
                if x < session.tracking.shelter[0]+session.tracking.shelter[2]]:
                frames_to_shelter = [idx for idx,x in enumerate(session.tracking.x[t:-1])
                                     if x < session.tracking.shelter[0]+session.tracking.shelter[2]][0]
            else:
                continue

            if not [x for x in threat_out if s<x<s+frames_to_shelter]:
                last_entrance = [x for x in threat_in if x < t][-1]
                trans_t_to_s.append((t, t + frames_to_shelter))

    setattr(transitions, 'shelter_to_threat', trans_s_to_t)
    setattr(transitions, 'threat_to_shelter', trans_t_to_s)
    setattr(session, 'transitions', transitions)
    return session

class Transitions:
    def __init__(self):
        pass

def get_all_sessions_mouse_experiment(datalog, k, session):
    mouse = session.mouseid[0][0]
    exp = session.experiment[0][0]

    keys, sessions = [], []
    for k,sess in datalog.__dict__.items():
        if k[0] == 's':
            if sess.mouseid[0][0] == mouse and sess.experiment[0][0] == exp:
                keys.append(k)
                sessions.append(sess)

    return keys, sessions, [exp + mouse]

def three_arms_maze_velocity_analysis(session):
    # get vars
    vel = np.vstack([session.aligned.vs.velocity[:,idx] for idx,c in enumerate(session.trial_escapes.vs) if c]).T
    x = np.vstack([session.aligned.vs.x[:,idx] for idx,c in enumerate(session.trial_escapes.vs) if c]).T
    y = np.vstack([session.aligned.vs.y[:,idx] for idx,c in enumerate(session.trial_escapes.vs) if c]).T
    direction = np.vstack([session.aligned.vs.direction[:,idx] for idx,c in enumerate(session.trial_escapes.vs) if c]).T

    trial_length = vel.shape[0]
    half_trial_length = round(trial_length/2)

    # Get vel_baseline
    baseline_start = round(half_trial_length-3*session.trackingsect0.fps)
    bl_std = np.std(vel[baseline_start:half_trial_length-1])
    bl_length = half_trial_length-baseline_start
    vel_baseline = np.mean(vel[baseline_start:half_trial_length-1])

    # Now plot
    pre_pad, post_pad = round(5*session.trackingsect0.fps), round(7.5*session.trackingsect0.fps)
    ax_len = pre_pad+post_pad
    bl_start = ax_len-bl_length
    bl_start = 1-(bl_start/ax_len)
    bl_end=pre_pad/ax_len

    f, axarr = plt.subplots(2,2, gridspec_kw={'height_ratios':[3, 1], 'width_ratios':[2, 1]})
    ax = axarr[0][0]
    ax2 = axarr[1][0]
    ax3 = axarr[0][1]
    ax4 = axarr[1][1]

    # Plot individual trials
    for i in range(vel.shape[1]):
        ax.plot(plotting_functions.line_smoother(vel[half_trial_length-pre_pad:half_trial_length+post_pad, i], 51,3),
                c='0.7', lw=0.75, alpha=0.5)

    # plot mean
    smoothed_mean = plotting_functions.line_smoother(np.mean(vel, axis=1)[half_trial_length-pre_pad:half_trial_length+post_pad], 51, 3)
    ax.plot(smoothed_mean, c='r', lw=1.5, alpha=0.75)

    # plot vel_baseline and std
    ax.axhline(xmin=bl_start, xmax=bl_end, y=vel_baseline, c='g')
    ax.axhline(xmin=bl_start, xmax=bl_end, y=vel_baseline-bl_std, c='g', lw=0.75, ls='--')
    ax.axhline(xmin=bl_start, xmax=bl_end, y=vel_baseline+bl_std, c='g', lw=0.75, ls='--')

    # mark stim onset
    ax.axvline(pre_pad+1, color='k', ls='--', lw=1)
    ax.set_title('velocity')

    # plot derivatives of mean velocity
    ax2.plot(plotting_functions.line_smoother(np.diff(smoothed_mean), 51, 3), c='g', lw=1.5, alpha=0.76)
    ax2.plot(plotting_functions.line_smoother(np.diff(smoothed_mean, 2), 51, 3), c='b', lw=1.5, alpha=0.76)
    ax2.axvline(pre_pad+1, color='k', ls='--', lw=1)
    ax2.set_title('acceleration')

    # Plot X and Y movements
    ax3.plot(y[half_trial_length-pre_pad:half_trial_length+post_pad], c='g', lw=1, alpha=0.25)
    ax3.plot(np.mean(y, axis=1)[half_trial_length-pre_pad:half_trial_length+post_pad], c='g', lw=2, alpha=0.76)
    ax3.set_title('Y displacement')

    # Plot orientation
    ax4.plot(direction[half_trial_length-pre_pad:half_trial_length+post_pad], c='g', lw=1, alpha=0.25)
    ax4.plot(np.mean(direction, axis=1)[half_trial_length-pre_pad:half_trial_length+post_pad], c='g', lw=2, alpha=0.76)
    ax4.set_title('Y displacement')

    plt.suptitle('{}-{}'.format(session.date[0][0], session.mouseid[0][0]))

    a = 1

def three_arms_maze(session, display=False):
    # Figure out from which the mouse came from and along which it fled for each stim
    wnd = round(session.aligned.vs.x.shape[0]/2)

    # add here the possibility to use user-defined ROI to set these boundaries
    outcomes = []
    for idx in range(session.aligned.vs.x.shape[1]):
        # check if there is a second stimulus within a minute, if this is the case: two stimuli where delivered
        # one after the other
        if session.trial_escapes.vs[idx]:
            # Check when the mouse is midway through the short arm range to check escape/origin routs
            stim_pos = (session.aligned.vs.x[wnd, idx], session.aligned.vs.y[wnd, idx])

            s_bridge = session.short_bridge
            # Get time at which mouse leaves and reaches a Y location equivalent to the midpoint of the S bridge
            y_in_range_pre = np.where(session.aligned.vs.y[0:wnd, idx]<s_bridge[1]+(0.5*s_bridge[3]))[-1][-1]
            y_in_range_post = np.where(session.aligned.vs.y[wnd+1:-1, idx]<s_bridge[1]+(0.5*s_bridge[3]))[-1][0]-1

            # Check where it is on the X axis at those time points
            pre_pos = session.aligned.vs.x[y_in_range_pre, idx]
            post_pos = session.aligned.vs.x[wnd+y_in_range_post, idx]

            # Determine escape and orgine routes
            if pre_pos < s_bridge[0]:  # coming from L bridge
                pre = 1
            elif pre_pos > s_bridge[2]+s_bridge[0]:  # coming from M bridge
                pre = 3
            else:
                pre = 2  # coming from S bridge

            if post_pos < s_bridge[0]-s_bridge[0]/4:
                post = 1
            elif post_pos > s_bridge[2]+s_bridge[0]+s_bridge[0]/4:
                post = 3
            else:
                post = 2

            # Store results
            outcomes.append((pre, post))

    # Calculate escape probabilities for each trial type (based on arm of origin) for each escape possibility
    probs = np.zeros((3,3))
    from_long = [x for x in outcomes if x[0]==1]
    from_short = [x for x in outcomes if x[0]==2]
    from_medium = [x for x in outcomes if x[0]==3]

    # in probs: 0 --> L, 1 --> S, 2 --> M
    # such that probs[0, 1] --> probs from L to S

    if from_long:
        probs[0, 0] = len([x for x in from_long if x[1]==1])/len(from_long)
        probs[0, 1] = len([x for x in from_long if x[1]==2])/len(from_long)
        probs[0, 2] = len([x for x in from_long if x[1]==3])/len(from_long)

    if from_short:
        probs[1, 0] = len([x for x in from_short if x[1]==1])/len(from_short)
        probs[1, 1] = len([x for x in from_short if x[1]==2])/len(from_short)
        probs[1, 2] = len([x for x in from_short if x[1]==3])/len(from_short)

    if from_medium:
        probs[2, 0] = len([x for x in from_medium if x[1]==1])/len(from_medium)
        probs[2, 1] = len([x for x in from_medium if x[1]==2])/len(from_medium)
        probs[2, 2] = len([x for x in from_medium if x[1]==3])/len(from_medium)

    # Get counts of escapes to and from each arm
    count_from_l = len(from_long)
    count_from_s = len(from_short)
    count_from_m = len(from_medium)
    all_counts_origin = [count_from_l, count_from_s, count_from_m]

    all_counts_escapes = [0, 0, 0]
    for idx, count in enumerate(all_counts_origin):
        all_counts_escapes[0] += int(round(probs[idx, 0] * count))
        all_counts_escapes[1] += int(round(probs[idx, 1] * count))
        all_counts_escapes[2] += int(round(probs[idx, 2] * count))

    if session.mouseid[0][0] == 'CA2765':
        a = 1

    # Display results
    if display:
        f, ax = plt.subplots()
        xlabels = ['', 'Long - {}'.format(all_counts_escapes[0]), 'Short - {}'.format(all_counts_escapes[1]),
                   'Medium - {}'.format(all_counts_escapes[2])]
        ylabels = ['', 'Long \n {}'.format(count_from_l), 'Short \n {}'.format(count_from_s),
                   'Medium \n {}'.format(count_from_m)]
        ax.matshow(probs, cmap='bwr')

        for i in range(3):
            for j in range(3):
                ax.text(j, i, '{}'.format(round(probs[i, j],2)), ha='center', va='center',backgroundcolor='white',
                                          bbox=dict(boxstyle='round', facecolor='white', edgecolor='0.3'))

        ax.set_xticklabels(xlabels)
        ax.set_yticklabels(ylabels)
        ax.set_yticks(np.arange(-.5, 3, 1), minor=True)
        ax.grid(which='minor', color='k', linestyle='-', linewidth=2)
        ax.set_title('Transition probability - {} - {}'. format(session.date[0][0], session.mouseid[0][0]))

    return probs, [from_long, from_short, from_medium]


from collections import Counter
import argparse
import string
import os

def load_data_from_file(filepath):
    with open(filepath, 'r') as file_object:
        str_data = file_object.read()
    return str_data

def get_most_frequent_words(text, output_word_count):
    letter_text = text.lower()
    for symbol in string.punctuation:
        letter_text.replace(symbol, ' ')
    word_list = letter_text.split()
    word_statistic = Counter(word_list).most_common(output_word_count)
    return word_statistic

def check_is_natural(input_string, default_value=10):
    if not input_string.isdigit():
        return default_value
    if int(input_string) > 0:
        return int(input_string)
    else:
        return default_value

def parse_arguments():
    parser = argparse.ArgumentParser(
        description='The word frequency counter.'
    )
    parser.add_argument(
        'filepath',
        type=str,
        help='a path to the text file where the script counts word frequency'
    )
    parser.add_argument(
        '-c',
        '--output_word_count',
        type=check_is_natural,
        default=10,
        help='how many words with statistic print as a result'
    )
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_arguments()
    if not os.path.exists(args.filepath):
        exit("File doesn't exist")
    text = load_data_from_file(args.filepath)
    if text is '':
        exit("The file is empty")
    for word, frequency in get_most_frequent_words(text, args.output_word_count):
        print('{word} - {frequency}'.format(word=word, frequency=frequency))


# Initial version, just pulls timestamps and only writes to csv after completes

import requests
import pdb
import time
import pandas as pd

index = 201910250000
urls = []
failed_indices = []
while index < 201912022100:
#while index < 201910262100:
    response = requests.get("https://archive.org/wayback/available?url=teamtrees.org&timestamp=" + str(index))
    try:
        current_url = response.json()['archived_snapshots']['closest']['url']
        urls.append(current_url)
        index = index + 1
    except:
        print(str(index) + ' failed, waiting 1 minute')
        time.sleep(30)

    if str(index)[10:] == '60':
        index = index - 60 + 100
    if str(index)[8:] == '2400':
        index = index - 2400 + 10000
    if str(index)[6:8] == '31':
        if str(index)[4:6] == '11':
            index = 201912010000
    elif str(index)[6:8] == '32':
        if str(index)[4:6] == '10':
            index = 201911010000
    print(index)

urls = pd.DataFrame(urls)
urls = urls.drop_duplicates()
urls.to_csv('urls_1min.csv',mode='a')
pdb.set_trace()


import argparse
import logging

from context import Context

def parse_args(args, os_family):
    parser = argparse.ArgumentParser(
        description="Script for building genny",
        epilog="Unknown positional arguments will be forwarded verbatim to the cmake"
        " invocation where relevant",
    )

    # Python can't natively check the distros of our supported platforms.
    # See https://bugs.python.org/issue18872 for more info.
    parser.add_argument(
        "-d",
        "--linux-distro",
        choices=["ubuntu1804", "archlinux", "rhel8", "rhel70", "rhel62", "amazon2", "not-linux"],
        help="specify the linux distro you're on; if your system isn't available,"
        " please contact us at #workload-generation",
    )
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument(
        "-g", "--run-global", action="store_true", help="allow installation outside of a virtualenv"
    )
    parser.add_argument(
        "-i",
        "--ignore-toolchain-version",
        action="store_true",
        help="ignore the toolchain version, useful for testing toolchain changes",
    )
    parser.add_argument(
        "-b",
        "--build-system",
        choices=["make", "ninja"],
        default="ninja",
        help="Which build-system to use for compilation. May need to use make for " "IDEs.",
    )
    parser.add_argument("-s", "--sanitizer", choices=["asan", "tsan", "ubsan"])

    subparsers = parser.add_subparsers(
        dest="subcommand",
        description="subcommands perform specific actions; make sure you run this script without "
        "any subcommand first to initialize the environment",
    )
    subparsers.add_parser(
        "cmake-test", help="run cmake unit tests that don't connect to a MongoDB cluster"
    )
    subparsers.add_parser("benchmark-test", help="run benchmark unit tests")

    resmoke_test_parser = subparsers.add_parser(
        "resmoke-test", help="run cmake unit tests that connect to a MongoDB cluster"
    )
    group = resmoke_test_parser.add_mutually_exclusive_group()
    group.add_argument(
        "--suites", dest="resmoke_suites", help='equivalent to resmoke.py\'s "--suites" option'
    )
    group.add_argument(
        "--create-new-actor-test-suite",
        action="store_true",
        dest="resmoke_cnats",
        help='Run the "genny_create_new_actor" resmoke test suite,'
        " incompatible with the --suites options",
    )
    resmoke_test_parser.add_argument(
        "--mongo-dir",
        dest="resmoke_mongo_dir",
        help="path to the mongo repo, which contains buildscripts/resmoke.py",
    )

    subparsers.add_parser("install", help="just run the install step for genny")
    subparsers.add_parser("clean", help="cleanup existing build")
    subparsers.add_parser("self-test", help="run lamplib unittests")

    known_args, unknown_args = parser.parse_known_args(args)

    if os_family == "Linux" and not known_args.subcommand and not known_args.linux_distro:
        raise ValueError("--linux-distro must be specified on Linux")

    return known_args, unknown_args

def add_args_to_context(args):
    """
    Add command line arguments to the global context object to be used later on.

    Consider putting command line arguments onto the context if it is used by more
    than one caller.

    :param args:
    :return:
    """
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
    Context.IGNORE_TOOLCHAIN_VERSION = args.ignore_toolchain_version
    Context.BUILD_SYSTEM = args.build_system
    Context.SANITIZER = args.sanitizer


from utils.data_loader import prepare_data_seq
from utils import config
from model.transformer import Transformer
from model.transformer_mulexpert import Transformer_experts
from model.common_layer import evaluate, count_parameters, make_infinite
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import xavier_uniform_
from copy import deepcopy
from tqdm import tqdm
import os
import time
import numpy as np
import math
from tensorboardX import SummaryWriter
from utils.data_loader import collate_fn
import pickle
from utils.data_loader import *
from utils.beam_omt_experts import Translator
import copy

#from torchsummary import summary

moel_model = torch.load("saved_models_testing/saved_model60000.pt")
#moel_model = moel_model.cuda()

data_loader_tra, data_loader_val, data_loader_tst, vocab, program_number = prepare_data_seq(batch_size=config.batch_size)
#print(data_loader_tst)
#evaluate(moel_model,data_loader_tst,ty="test", max_dec_step=50)

print("Value of config.label_smoothing : ",config.label_smoothing)
print("Value of naom : ",config.noam)
print("Vaue of cuda :",config.USE_CUDA)
print("Value of pretrain_emb : ",config.pretrain_emb)
print("Value of softmax : ",config.softmax)
print("Value of basic_listener : ",config.basic_learner)
print("Value of save_path : ",config.save_path)

# s = "i am sad"
s = ""
dialouge_conv = []
situation = []
data = {}
print("Enter the emotion of the conversation : ")
Emotion = input()
data['emotion'] = [Emotion]
#data['target'] = [["why","are","you","sad","?"]]
data['situation'] = []
data['context'] = []
if(os.path.exists('empathetic-dialogue/dataset_preproc.p')):
       print("LOADING empathetic_dialogue")
       with open('empathetic-dialogue/dataset_preproc.p', "rb") as f:
           [data_tra, data_val, data_tst, vocab] = pickle.load(f)
logging.info("Vocab  {} ".format(vocab.n_words))
t = Translator(moel_model, vocab)
print("start")
s = input()
# # sent11=["I","am","sad"]
# sent11 = s.split(' ')
# sent22=["I","am","happy"]
# conv1=[sent11]
# conv2=[sent11]
# conv2.append(sent22)
# diag=[conv1]
# diag.append(conv2)
# print(diag)
#diag=[]
conv = []
while(s!="exit"):
    # print(dialouge)
    if(conv == []):
        conv1 = conv
        conv1 = [s.split(" ")]
    else:
        conv1 = conv
        conv1.append(s.split(" "))
    data['context'].append(conv1)
    conv = copy.deepcopy(conv1)
    data["emotion"].append(data["emotion"][0])
    # data['context'].append(dialouge_conv)
    # print(len(data['context']))
    # for i in range(len(data['context'])):
    #     print(i,data['context'][i])
    if(situation==[]):
        situation = s.split(' ')
    data['situation'].append(situation)
    if(len(data["context"])==3):
        data["context"] = data["context"][1:3]
        data["emotion"] = data["emotion"][1:3]
        data["target"] = data["target"][1:3]
        data["situation"] = data["situation"][1:3]
    # print(data)
    # data['context'][0]= [data_loader['context'][0][0]]
    # print("before")
    # print(data['context'])
    dataset_input = Dataset(data, vocab)
    # print(dataset_input)
    data_loader = torch.utils.data.DataLoader(dataset=dataset_input,
                                        batch_size=32,
                                        shuffle=False, collate_fn=collate_fn)
    # print("after")
    # print(data['context'])
    # s = input()
    # print(data['context'])
    # data_loader = data_loader.gpu()

    #change=================19-02-2020

    #pbar = tqdm(enumerate(data_loader),total=len(data_loader))

    # print("Dialogue CONV",dialouge_conv)
    # print("data[context] ",data["context"])
    #pbar = tqdm(enumerate(data_loader))
    #for j, batch in pbar:
    #    batch["input_batch"]=batch["input_batch"].cuda()
    #    batch["input_lengths"]=batch["input_lengths"].cuda()
    #    batch["mask_input"]=batch["mask_input"].cuda()
    #    batch["target_batch"]=batch["target_batch"].cuda()
    #    batch["target_lengths"]=batch["target_lengths"].cuda()
    #    #print("Here")
    #    #print(batch)
    #    sent_g = moel_model.decoder_greedy(batch,max_dec_step=60)
    #    sent_b = t.beam_search(batch, max_dec_step=60)
    #    sent_t = moel_model.decoder_topk(batch, max_dec_step=60)
    #    print("Greedy : ",sent_g)
    #    print("Beam : ",sent_b)
    #    print("Output from topk : ",sent_t)
    #    try:
    #        data["target"].append(sent_t[0].split(' '))
    #    except Exception:
    #        data["target"] = [sent_t[0].split(' ')]
    #    dialouge_conv.append(sent_t[0].split(' '))
        # data['emotion'].append(data["emotion"][0])
        # data['emotion'].append(data["emotion"][0])
        #print(sent_g)
    #=======commented---till---here---==================
    sent_g = model.decoder_greedy(batch,max_dec_step=max_dec_step)
    sent_b = t.beam_search(batch, max_dec_step=max_dec_step)
    # print(sent_g,sent_b)
    sent_t = model.decoder_topk(batch, max_dec_step=max_dec_step)
    for i, (greedy_sent, beam_sent)  in enumerate(zip(sent_g, sent_b)):
        rf = " ".join(batch["target_txt"][i])
        hyp_g.append(greedy_sent)
        hyp_b.append(beam_sent)
        #hyp_t.append(topk_sent)
        ref.append(rf)
        print_custum(emotion= batch["program_txt"][i],
                    dial=[" ".join(s) for s in batch['input_txt'][i]] if config.dataset=="empathetic" else " ".join(batch['input_txt'][i]),
                    ref=rf,
                    #hyp_t=topk_sent,
                    hyp_g=greedy_sent,
                    hyp_b=beam_sent)
    s= input()

    #break


"""
Entity schema validation
"""

# pylint: disable=C0103

# stdlib
import re
from typing import Callable

# library
from voluptuous import All, Invalid, Length

def MatchesRE(name: str, pattern: str) -> Callable:
    """
    Returns a validation function that checks if a string matches a regex pattern
    """
    expr = re.compile(pattern)

    def mre(txt: str) -> str:
        """
        Raises an exception if a string doesn't match the required format
        """
        if expr.fullmatch(txt) is None:
            raise Invalid(f"'{txt}' is not a valid {name}")
        return txt

    return mre

Token = All(
    str,
    Length(min=10),
    lambda x: x.strip().split()[-1],
    MatchesRE("token", r"[A-Za-z0-9\-\_]+"),
)