import os
import subprocess
import tempfile
from pathlib import Path
from typing import List

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from tqdm import tqdm

from metavideo import get_metagrid
from object_tracking_operations import (extract_masked_object_clips,
                                        extract_obj_gifs_parallel,
                                        extract_object_thumbs,
                                        interpolate_missing_data,
                                        merge_with_chromakey)
from utils import ensure_coords, ensure_dir, find_video_by_id, uniquify

def select_shots_by_entity(annotation_data: pd.DataFrame,
                           key: List[str], threshold=0.5,
                           search_categories=False,
                           padding_before=0,
                           padding_after=0):
    """
    Select shots by key entity
    :param annotation_data: dataframe with annotation data
    :param key: key entity to filter by
    :param threshold: minimum confidence score to consider
    :param search_categories: search in categories as well
    :param padding_before: amount of seconds to subtract from start time
    :param padding_after: amount of seconds to add to end time
    :return: dataframe with selected shots
    """
    if type(key) == str:
        key = [key]

    if search_categories and 'category' not in annotation_data.columns:
        raise KeyError("The dataframe has no 'category' column")
    if 'entity' not in annotation_data.columns:
        raise KeyError("The dataframe has no 'entity' column")
    elif all(k not in annotation_data['entity'].unique() for k in key):
        raise ValueError(f"Key entity {key} not found in dataframe")

    entity_shots = annotation_data[annotation_data['entity'].str.lower().isin(key)].reset_index(drop=True)

    if search_categories:
        category_shots = annotation_data[annotation_data['category'].str.lower().isin(key)].reset_index(drop=True)
        selected_shots = pd.concat([entity_shots, category_shots]).drop_duplicates(keep='first')
    else:
        selected_shots = entity_shots

    selected_shots = selected_shots[selected_shots['confidence'] >= threshold].reset_index(drop=True)

    if padding_before:
        selected_shots['start_sec'] -= padding_before
        selected_shots['start_sec'] = selected_shots['start_sec'].apply(lambda x: max(x, 0))

    if padding_after:
        end_time = entity_shots['end_sec'].max()
        selected_shots['end_sec'] += padding_after
        selected_shots['end_sec'] = selected_shots['end_sec'].apply(lambda x: max(x, end_time))

    return selected_shots

def select_shots_by_keyword(annotation_data: pd.DataFrame,
                            key: List[str],
                            threshold=0.5,
                            padding_before=5,
                            padding_after=3):
    key = [key.lower()] if type(key) == str else [k.lower() for k in key]
    if 'word' not in annotation_data.columns:
        raise KeyError("The dataframe has no 'word' column")
    elif all(k not in annotation_data['word'].unique() for k in key):
        raise ValueError(f"Keyword {key} not found in dataframe")

    selected_shots = annotation_data[annotation_data['word'].str.lower().str.replace(r'[^\w\s]+', '').isin(key)]
    selected_shots = selected_shots[selected_shots['confidence'] >= threshold]
    selected_shots = add_padding_shots(annotation_data, padding_after, padding_before, selected_shots)
    selected_shots = selected_shots.drop_duplicates(keep='first')
    selected_shots = selected_shots.sort_index()
    selected_shots = selected_shots.reset_index(drop=True)

    starts = selected_shots['start_sec'].tolist()
    ends = selected_shots['end_sec'].tolist()

    selected_shots = merge_consecutive(selected_shots)

    return selected_shots

def select_shots_by_consecutive_words(annotation_data: pd.DataFrame,
                                      key: List[str]):
    key = [k.lower().strip() for k in key] if isinstance(key, list) else [k.lower().strip() for k in key.split(',')]
    annotation_data['word'] = annotation_data['word'].str.lower().str.replace(r'[^\w\s]+', '').str.strip()
    selected_shots = pd.DataFrame()
    indexes = annotation_data[annotation_data['word'].isin(key)].index.values

    # for each index in indexes, get the following n indexes, where n is the length of the key
    for i in indexes:
        following_indexes = [i + j for j in range(len(key))]
        # check that the 'word' column of the first row in the following_indexes list matches the first word in the key
        # that the second row in the following_indexes list matches the second word in the key, etc.
        if all(annotation_data.iloc[m]['word'] == key[n] for n, m in enumerate(following_indexes)):
            # if all the words in the key match the words in the following_indexes list, append a new row to the
            # selected_shots dataframe with the start and end times of the first and last rows in the following_indexes
            # list, and the words in the key
            selected_shots = selected_shots.append({
                'id': annotation_data.iloc[following_indexes[0]]['id'],
                'word': ' '.join(key),
                'start_sec': annotation_data.iloc[following_indexes[0]]['start_sec'],
                'end_sec': annotation_data.iloc[following_indexes[-1]]['end_sec'],
                'id0': following_indexes[0],
                'id1': following_indexes[-1]
            }, ignore_index=True)

            # selected_shots = selected_shots.append(annotation_data.iloc[following_indexes])

    # remove duplicates and sort by index
    selected_shots = selected_shots.drop_duplicates(keep='first')
    selected_shots = selected_shots.sort_index()
    return selected_shots

def merge_consecutive(selected_shots):
    prev_len = len(selected_shots)
    while True:
        for i, row in selected_shots.iterrows():
            if i < len(selected_shots) - 1:
                current_end = float(row['end_sec'])
                next_start = float(selected_shots.iloc[i + 1]['start_sec'])
                if 0 <= next_start - current_end <= 2:
                    selected_shots.loc[i, "end_sec"] = selected_shots.iloc[i + 1]["end_sec"]
                    selected_shots.loc[i, "word"] = selected_shots.iloc[i]["word"] + " " + selected_shots.iloc[i + 1][
                        "word"]
                    selected_shots = selected_shots.drop(i + 1)
                    selected_shots = selected_shots.reset_index(drop=True)
        current_len = len(selected_shots)
        print("current_len: ", current_len, end="\r")
        if current_len < prev_len:
            prev_len = current_len
        else:
            break
    return selected_shots

def add_padding_shots(annotation_data, padding_after, padding_before, selected_shots):
    for n, shot in selected_shots.iterrows():
        if n == 0:
            continue
        for i in range(1, padding_before + 1):
            if n - i < 0:
                break
            current_start = int(shot['start_sec'])
            previous_end = int(annotation_data.iloc[n - i]['end_sec'])
            if current_start - previous_end <= 2:
                selected_shots = selected_shots.append(annotation_data.iloc[n - i])
            else:
                break
        for i in range(1, padding_after + 1):
            if n + i > len(annotation_data):
                break
            current_end = int(shot['end_sec'])
            next_start = int(annotation_data.iloc[n + i]['start_sec'])
            diff = next_start - current_end
            if next_start - current_end <= 2:
                selected_shots = selected_shots.append(annotation_data.iloc[n + i])
            else:
                break
    return selected_shots

def add_padding_to_consecutive_keywords(annotation_data, padding_after, padding_before, selected_shots):
    for n, shot in selected_shots.iterrows():
        if n == 0:
            continue
        for i in range(1, padding_before + 1):
            if n - i < 0:
                break
            current_start = int(shot['start_sec'])
            previous_end = int(annotation_data.iloc[shot['id0'] - i]['end_sec'])
            if current_start - previous_end <= 2:
                selected_shots = selected_shots.append(annotation_data.iloc[shot['id0'] - i])
            else:
                break
        for i in range(1, padding_after + 1):
            if n + i > len(annotation_data):
                break
            current_end = int(shot['end_sec'])
            next_start = int(annotation_data.iloc[shot['id1'] + i]['start_sec'])
            diff = next_start - current_end
            if next_start - current_end <= 2:
                selected_shots = selected_shots.append(annotation_data.iloc[shot['id1'] + i])
            else:
                break
        print(selected_shots)
    return selected_shots

def extract_shots(_df: pd.DataFrame, in_dir: Path, out_dir: Path, text: str = False):
    """
    :param text:
    :param out_dir:
    :param in_dir:
    :param _df: DataFrame with data on which shots to extract
    :return: None
    """
    df = _df.sort_values(by=['id', 'start_sec'])
    out_dir.mkdir(parents=True, exist_ok=True)

    for index, row in tqdm(df.iterrows(), total=df.shape[0], desc='Extracting shots'):
        entity = row[text] if text else row[0]
        video_id = row['id']
        filename = Path(in_dir, find_video_by_id(row['id'], in_dir))
        start = "%.2f" % row['start_sec']
        end = "%.2f" % row['end_sec']
        if start == end:
            continue
        in_path = filename.resolve().as_posix()
        out_path = uniquify(os.path.join(out_dir.as_posix(), f"{video_id}.mp4"))
        command = ["ffmpeg"]
        options = ["-i", in_path, "-ss", start, "-to", end, "-y", "-movflags",
                   "faststart", "-avoid_negative_ts", "1", "-acodec", "copy", out_path]

        if text:
            text_filter = ["drawtext=", "fontfile=Inter-Regular.ttf:", f"text='{entity}':",
                           "fontcolor=white:", "fontsize=24:", "box=1:", "boxcolor=black@0.5:", "boxborderw=5:",
                           "x=(w-text_w)/2:", "y=24"]
            filter_args = "".join(text_filter)
            options.insert(6, "-vf")
            options.insert(7, filter_args)

        args = command + options
        operation = subprocess.run(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if operation.returncode != 0:
            print(operation.stderr)
            raise RuntimeError("Ué uagliù è succiesso nu guaio mentre stev tagliann e' video, liv 'a miezz "
                               "'stderr=subprocess.DEVNULL' e vir nu poc ch'è succiess")

# TODO add threading and hardware acceleration because this can get pretty long and boring
def merge_shots(in_dir: Path, out_dir: Path):
    """
    Merge shots into one video
    :param selected_shots_path: path to selected shots
    :return: None
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    files = [file.as_posix() for file in in_dir.glob('*.mp4')]
    # select only files have an audio and a video stream
    files = [file for file in files if len(subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries", "stream=codec_type", "-of",
         "default=noprint_wrappers=1:nokey=1", file],
        stdout=subprocess.PIPE).stdout.decode('utf-8').splitlines()) > 0 and len(subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=codec_type", "-of",
         "default=noprint_wrappers=1:nokey=1", file],
        stdout=subprocess.PIPE).stdout.decode('utf-8').splitlines()) > 0]
    files.sort()

    out_path = uniquify(Path(out_dir, 'merged.mp4').as_posix())
    command = ["ffmpeg"]
    input_files = [["-i", file] for file in files]
    input_files = [item for sublist in input_files for item in sublist]
    streams = [f"[{i}:v][{i}:a]" for i in range(len(files))]
    concat = [f"concat=n={len(files)}:v=1:a=1[v][a]"]
    mapper = ['-map', '[v]', '-map', '[a]']
    sync = ["-vsync", "2", '-threads', '0']
    options = input_files + ["-filter_complex"] + [f"{''.join(streams + concat)}"] + mapper + sync + [out_path]
    args = command + options
    print(args)
    operation = subprocess.run(args, stdout=subprocess.DEVNULL)

def agnostic_merge(video_dir, output_dir):
    output_directory = ensure_dir(output_dir)
    files = [os.path.join(video_dir, file) for file in os.listdir(video_dir) if file.endswith('.mp4')]
    out_path = uniquify(os.path.join(output_directory, 'merged.mp4'))
    command = ["ffmpeg"]
    input_files = [["-i", file] for file in files]
    input_files = [item for sublist in input_files for item in sublist]
    aspect_ratio_handler = [
        f'[{i}]scale=1280:720:force_original_aspect_ratio=decrease,pad=1280:720:(ow-iw)/2:(oh-ih)/2,setsar=1[v{i}];' for
        i in range(len(files))]
    streams = [f'[v{i}][{i}:a:0]' for i in range(len(files))]
    concat = [f"concat=n={len(files)}:v=1:a=1[v][a]"]
    mapper = ['-map', '[v]', '-map', '[a]']
    sync = ["-vsync", "2"]
    options = input_files + ["-filter_complex"] + [
        f"{''.join(aspect_ratio_handler + streams + concat)}"] + mapper + sync + [out_path]
    args = command + options
    subprocess.run(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def render_heatmap(out_dir: Path, data: pd.DataFrame, key: list, **kwargs) -> None:
    """
    :param data: DataFrame with data to plot
    :param out_dir: folder where to save the heatmap
    :keyword resolution: tuple(width, height)
    :return: None
    """
    print(data)

    data = data[data['object_name'].isin(key)]
    print(data)
    res = (1920, 1080) if kwargs.get('resolution') is None else kwargs.get('resolution')
    data.fillna(0, inplace=True)

    img = np.zeros((res[1], res[0], 4), dtype=np.uint)
    for _, row in data.iterrows():
        left, top, right, bottom = ensure_coords(row['left'], row['top'], row['right'], row['bottom'])
        img[int(top * res[1]):int(bottom * res[1]), int(left * res[0]):int(right * res[0]), 0:3] += 1
    img = img / img.max()
    img[0:res[1], 0:res[0], 3] = 1

    from scipy.ndimage.filters import gaussian_filter
    img = gaussian_filter(img, sigma=5)

    from matplotlib.colors import LinearSegmentedColormap
    cmap = LinearSegmentedColormap.from_list("", ["blue", "red"])
    video_id, object_name = data['id'].iat[0], data['object_name'].iat[0]
    out_path = out_dir / f"{video_id}_{object_name}.png"
    plt.imshow(img, cmap=cmap)
    plt.imsave(uniquify(out_path.as_posix()), img, cmap=cmap)

def berensham(x1, y1, x2, y2):
    dx = abs(x2 - x1)
    dy = abs(y2 - y1)
    sx = 1 if x1 < x2 else -1
    sy = 1 if y1 < y2 else -1
    err = dx - dy
    while True:
        yield x1, y1
        if x1 == x2 and y1 == y2:
            break
        e2 = 2 * err
        if e2 > -dy:
            err -= dy
            x1 += sx
        if e2 < dx:
            err += dx
            y1 += sy

def render_traces(out_dir: Path, data: pd.DataFrame, key: list, **kwargs) -> None:
    """
    :param data: DataFrame with data to plot
    :param out_dir: folder where to save the traces
    :keyword resolution: tuple(width, height)
    :return: None
    """
    print(list(data))
    data = data[data['object_name'].isin(key)]
    res = (1920, 1080) if kwargs.get('resolution') is None else kwargs.get('resolution')
    data.fillna(0, inplace=True)
    img = np.zeros((res[1], res[0], 4), dtype=np.uint)

    # we need to cluster data by obj_id in order to follow the same object troughout the video
    data = data.groupby('object_id').apply(lambda x: x.sort_values('time_seconds'))

    # since centroids's positions are not guaranteed to be continuous, we need tro draw a line from point 'n' to point 'n+1'
    # we will use Bresenham's line algorithm to do so
    for _, group in data.groupby('object_id'):
        for i in range(len(group)-1):
            left, top, right, bottom = ensure_coords(group['left'].iloc[i], group['top'].iloc[i], group['right'].iloc[i], group['bottom'].iloc[i])
            centroid = (int((left + right) / 2 * res[0]), int((top + bottom) / 2 * res[1]))
            n_left, n_top, n_right, n_bottom = ensure_coords(group['left'].iloc[i+1], group['top'].iloc[i+1], group['right'].iloc[i+1], group['bottom'].iloc[i+1])
            next_centroid = (int((n_left + n_right) / 2 * res[0]), int((n_top + n_bottom) / 2 * res[1]))
            for x, y in berensham(centroid[0], centroid[1], next_centroid[0], next_centroid[1]):
                img[(y-5):(y+5), (x-5):(x), 0:3] += 1

    img = img / img.max()
    img[0:res[1], 0:res[0], 3] = 1

    video_id, object_name = data['id'].iat[0], data['object_name'].iat[0]
    out_path = out_dir / f"{video_id}_{object_name}.png"
    plt.imsave(uniquify(out_path.as_posix()), img)
    print(f"Rendered trace for {video_id} {object_name}")
    return None

def extract_object_thumbnails(in_dir: Path, out_dir: Path, data: pd.DataFrame, key: list) -> None:
    """
    Given a video and a dataframe with object tracking annotations extracts thumbnails of selected the objects
    :rtype: None
    :param key: name of the object to extract
    :param in_dir: directory where the source video is stored
    :param out_dir: directory where to save the thumbnails
    :param data: dataframe with object tracking annotations
    :return: None
    """
    data = data[data['object_name'].isin(key)]
    data = interpolate_missing_data(data)
    extract_object_thumbs(in_dir, out_dir, data)
    return None

def extract_object_gifs(in_dir: Path, out_dir: Path, data: pd.DataFrame, key: list) -> None:
    """
    Given a video and a dataframe with object tracking annotations extracts gifs of selected the objects
    :rtype: None
    :param key: name of the object to extract
    :param in_dir: directory where the source video is stored
    :param out_dir: directory where to save the thumbnails
    :param data: dataframe with object tracking annotations
    :return: None
    """
    data = data[data['object_name'].isin(key)]
    data = interpolate_missing_data(data)
    extract_obj_gifs_parallel(in_dir, out_dir, data)
    return None

def extract_masked_clips(in_dir: Path, out_dir: Path, data: pd.DataFrame, key: list) -> None:
    """
    Given a video and a dataframe with object tracking annotations extracts clips of selected the objects
    :param key: name of the object to extract
    :param in_dir: directory where the source video is stored
    :param out_dir: directory where to save the clips
    :param data: dataframe with object tracking annotations
    :return: None
    """
    data = data[data['object_name'].isin(key)]
    data = interpolate_missing_data(data)
    extract_masked_object_clips(in_dir, out_dir, data)

def extract_object_metavideo(in_dir: Path, out_dir: Path, data: pd.DataFrame, key: list) -> None:
    """
    Given a video and a dataframe with object tracking annotations extracts clips of selected the objects
    :param key: name of the object to extract
    :param in_dir: directory where the source video is stored
    :param out_dir: directory where to save the clips
    :param data: dataframe with object tracking annotations
    :return: None
    """
    data = data[data['object_name'].isin(key)]
    data = interpolate_missing_data(data)
    temp_dir = tempfile.TemporaryDirectory()
    extract_masked_object_clips(in_dir, Path(temp_dir.name), data, color=(0, 255, 0, 0))
    merge_with_chromakey(Path(temp_dir.name), out_dir)

def extract_object_metagrid(in_dir: Path, out_dir: Path, data: pd.DataFrame, key: list) -> None:
    """
    Given a video and a dataframe with object tracking annotations extracts clips of selected the objects
    :param key: name of the object to extract
    :param in_dir: directory where the source video is stored
    :param out_dir: directory where to save the clips
    :param data: dataframe with object tracking annotations
    :return: None
    """
    data = data[data['object_name'].isin(key)]
    data = interpolate_missing_data(data)
    temp_dir = tempfile.TemporaryDirectory()
    extract_object_thumbs(in_dir, Path(temp_dir.name), data)
    get_metagrid(Path(temp_dir.name), out_dir)


from __future__ import annotations

import configparser
import datetime
import os
import typing

import fabric
import invocations.console
import invoke

import invoke_patch
import src.constants

class GlobalConfig:
    host: str
    user: str
    key_filename: str
    project_name: str
    project_path: str

    source_filenames = [
        'main.py',
        'database.py',
        'utils.py',
        'telegram_utils.py',
        'analytics.py',
        'constants.py',

        'custom_logger.py',

        'config.cfg'
    ]
    meta_filenames = [
        'pyproject.toml',
        'poetry.lock'
    ]
    source_directories = [
        'migrations'
    ]

    @classmethod
    def load(cls) -> None:
        try:
            fabfile_config = configparser.ConfigParser()

            fabfile_config.read('fabfile.cfg')

            cls.host = fabfile_config.get('Fabric', 'Host')
            cls.user = fabfile_config.get('Fabric', 'User')
            cls.key_filename = os.path.expanduser(fabfile_config.get('Fabric', 'KeyFilename'))
            cls.project_name = fabfile_config.get('Fabric', 'ProjectName')
            cls.project_path = fabfile_config.get('Fabric', 'ProjectPath')
        except configparser.Error as error:
            raise invoke.Exit(
                message=f'Config error: {error}',
                code=1
            )

invoke_patch.fix_annotations()
GlobalConfig.load()

@fabric.task
def configure(connection: fabric.Connection) -> None:
    connection.user = GlobalConfig.user
    connection.inline_ssh_env = True
    connection.connect_kwargs.key_filename = GlobalConfig.key_filename

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'command': 'The shell command to execute on the server', 'env': 'An optional dictionary with environment variables'})
def execute(connection: fabric.Connection, command: str, env: typing.Dict[str, str] = None) -> None:
    if not command:
        return

    connection.run(command, env=env)

@fabric.task(pre=[configure], hosts=[GlobalConfig.host])
def cleanup(connection: fabric.Connection) -> None:
    question = f'Are you sure you want to completely delete the project "{GlobalConfig.project_name}" from "{GlobalConfig.host}"?'

    if invocations.console.confirm(
        question=question,
        assume_yes=False
    ):
        execute(connection, f'rm -rf {GlobalConfig.project_name}')
        execute(connection, f'rm -rf {GlobalConfig.project_path}/{GlobalConfig.project_name}')

@fabric.task(pre=[configure, cleanup], hosts=[GlobalConfig.host])
def setup(connection: fabric.Connection) -> None:
    execute(connection, f'mkdir -p {GlobalConfig.project_path}/{GlobalConfig.project_name}')
    execute(connection, f'ln -s {GlobalConfig.project_path}/{GlobalConfig.project_name} {GlobalConfig.project_name}')

    execute(connection, 'curl -sSL https://install.python-poetry.org | python -')

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'An optional filename to deploy to the server'})
def upload(connection: fabric.Connection, filename: typing.Optional[str] = None) -> None:
    def upload_file(file_format: str, file_name: str, destination_path_format='{.project_name}/{}') -> None:
        connection.put(file_format.format(file_name), destination_path_format.format(GlobalConfig, file_name))

    def upload_directory(directory_name: str) -> None:
        execute(connection, f'mkdir -p {GlobalConfig.project_name}/{directory_name}')

        for _root, _directories, files in os.walk(f'src/{directory_name}'):
            for file in files:
                upload_file(f'src/{directory_name}/{{}}', file, f'{{.project_name}}/{directory_name}/{{}}')

    if filename:
        if filename in GlobalConfig.source_directories:
            upload_directory(filename)
        else:
            if filename in GlobalConfig.source_filenames:
                file_path_format = 'src/{}'
            elif filename in GlobalConfig.meta_filenames:
                file_path_format = '{}'
            else:
                raise invoke.ParseError(f'Filename "{filename}" is not registered')

            upload_file(file_path_format, filename)
    else:
        for name in GlobalConfig.source_filenames:
            upload_file('src/{}', name)

        for name in GlobalConfig.meta_filenames:
            upload_file('{}', name)

        for directory in GlobalConfig.source_directories:
            upload_directory(directory)

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'An optional filename to deploy to the server'})
def deploy(connection: fabric.Connection, filename: typing.Optional[str] = None) -> None:
    upload(connection, filename)

    with connection.cd(GlobalConfig.project_name):
        execute(connection, 'eval "$(pyenv init --path)" && poetry install --no-dev', {
            'PATH': '$HOME/.pyenv/bin:$HOME/.poetry/bin:$PATH'
        })

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'The filename to backup locally from the server'})
def backup(connection: fabric.Connection, filename: str) -> None:
    current_date = datetime.datetime.now().strftime(src.constants.GENERIC_DATE_FORMAT)
    name, extension = os.path.splitext(filename)

    with connection.cd(GlobalConfig.project_name):
        connection.get(f'{GlobalConfig.project_name}/{filename}', f'backup_{name}_{current_date}{extension}')

@fabric.task(pre=[configure], hosts=[GlobalConfig.host])
def backup_db(context: fabric.Connection) -> None:
    backup(context, 'file_convert.sqlite')


import torch.nn as nn
import torch.multiprocessing as mp
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

import argparse
from lib.utils.file import bool_flag
from lib.utils.distributed import init_dist_node, init_dist_gpu, get_shared_folder
from lib.utils.flatten import flatten_dict

import submitit, random, sys
from pathlib import Path

def parse_args():
    parser = argparse.ArgumentParser(description='Template')
    parser.add_argument('--cfg', type=str, help='Configuration file')

    # PATHS
    paths_group = parser.add_argument_group("PATHS")
    paths_group.add_argument('--data', type=str, help='Path to raw dataset directory')
    paths_group.add_argument('--processed_data', type=str, help='Path to processed dataset directory')
    paths_group.add_argument('--processed_pattern', type=str, help='Pattern for processed data files')
    paths_group.add_argument('--genuine_pattern', type=str, help='Pattern for genuine data files')
    paths_group.add_argument('--simulated_pattern', type=str, help='Pattern for simulated data files')
    paths_group.add_argument('--out', type=str, help='Path to out directory')

    # GENERAL
    general_group = parser.add_argument_group("GENERAL")
    general_group.add_argument('--model', type=str, help='Model name')
    general_group.add_argument('--reset', action='store_true', help='Reset saved model logs and weights')
    general_group.add_argument('--tb', action='store_true', help='Start TensorBoard')
    general_group.add_argument('--gpus', type=str, help='GPUs list, only works if not on slurm')
    general_group.add_argument('--k_neighbors', type=int, help='Number of neighbors to consider in k-NN algorithm')

    # Model
    model_group = parser.add_argument_group("Model")
    model_group.add_argument('--num_layers', type=int, help='Number of layers')
    model_group.add_argument('--hidden_dim', type=int, help='Number of nodes in the hidden layer.')

    # Dataset
    dataset_group = parser.add_argument_group("Dataset")
    dataset_group.add_argument('--dataset', type=str, help='Dataset to choose')
    dataset_group.add_argument('--batch_per_gpu', type=int, help='Batch size per gpu')
    dataset_group.add_argument('--shuffle', action='store_true', help='Shuffle dataset')
    dataset_group.add_argument('--workers', type=int, help='Number of workers')
    dataset_group.add_argument('--keys', type=str, nargs='+', help='Columns to be used as data features')

    # Architecture
    architecture_group = parser.add_argument_group("Architecture")
    architecture_group.add_argument('--arch', type=str, help='Architecture to choose')

    # Trainer
    trainer_group = parser.add_argument_group("Trainer")
    trainer_group.add_argument('--trainer', type=str, help='Trainer to choose')
    trainer_group.add_argument('--epochs', type=int, help='Number of epochs')
    trainer_group.add_argument('--save_every', type=int, help='Save model every n epochs')
    trainer_group.add_argument('--fp16', action='store_true', help='Use fp16')

    # Optimization
    optimization_group = parser.add_argument_group("Optimization")
    optimization_group.add_argument('--optimizer', type=str, help='Optimizer to choose between "adam", "sgd", and "adagrad"')
    optimization_group.add_argument('--lr_start', type=float, help='Learning rate start')
    optimization_group.add_argument('--lr_end', type=float, help='Learning rate end')
    optimization_group.add_argument('--lr_warmup', type=int, help='Learning rate warmup')

    # SLURM
    slurm_group = parser.add_argument_group("SLURM")
    slurm_group.add_argument('--slurm', action='store_true', help='Use slurm')
    slurm_group.add_argument('--slurm_ngpus', type=int, help='Number of gpus per node')
    slurm_group.add_argument('--slurm_nnodes', type=int, help='Number of nodes')
    slurm_group.add_argument('--slurm_nodelist', type=str, help='Node list')
    slurm_group.add_argument('--slurm_partition', type=str, help='Partition')
    slurm_group.add_argument('--slurm_timeout', type=int, help='Timeout')

    args = parser.parse_args()

    # === Read CFG File === #
    if args.cfg:
        with open(args.cfg, 'r') as f:
            import ruamel.yaml as yaml
            yml = yaml.safe_load(f)

        # update values from cfg file only if not passed in cmdline
        cmd = [c[1:] for c in sys.argv if c[0]=='-']
        for k,v in flatten_dict(yml):
            if k not in cmd:
                args.__dict__[k] = v

    return args

class SLURM_Trainer(object):
    def __init__(self, args):
        self.args = args

    def __call__(self):

        init_dist_node(self.args)
        train(None, self.args)

def main():

    args = parse_args()
    args.port = random.randint(49152,65535)

    if args.slurm:

        # Almost copy-paste from https://github.com/facebookresearch/deit/blob/main/run_with_submitit.py
        args.output_dir = get_shared_folder(args) / "%j"
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
        executor = submitit.AutoExecutor(folder=args.output_dir, slurm_max_num_timeout=30)

        if hasattr(args, "slurm_account"):
            executor.update_parameters(
                mem_gb=12*args.slurm_ngpus,
                gpus_per_node=args.slurm_ngpus,
                tasks_per_node=args.slurm_ngpus,
                cpus_per_task=2,
                nodes=args.slurm_nnodes,
                timeout_min=2800,
                slurm_partition=args.slurm_partition,
                account=args.slurm_account
            )
        else:
            executor.update_parameters(
                mem_gb=12*args.slurm_ngpus,
                gpus_per_node=args.slurm_ngpus,
                tasks_per_node=args.slurm_ngpus,
                cpus_per_task=2,
                nodes=args.slurm_nnodes,
                timeout_min=2800,
                slurm_partition=args.slurm_partition
            )

        if args.slurm_nodelist:
            executor.update_parameters(slurm_additional_parameters = {"nodelist": f'{args.slurm_nodelist}' })

        executor.update_parameters(name=args.model)
        trainer = SLURM_Trainer(args)
        job = executor.submit(trainer)
        print(f"Submitted job_id: {job.job_id}")

    else:
        init_dist_node(args)
        mp.spawn(train, args = (args,), nprocs = args.ngpus_per_node)

def train(gpu, args):

    # === SET ENV === #
    init_dist_gpu(gpu, args)

    # === DATA === #
    get_dataset = getattr(__import__("lib.datasets.{}".format(args.dataset), fromlist=["get_dataset"]), "get_dataset")
    dataset = get_dataset(args)

    sampler = DistributedSampler(dataset, shuffle=args.shuffle, num_replicas = args.world_size, rank = args.rank, seed = 31)
    loader = DataLoader(dataset=dataset,
                        sampler = sampler,
                        batch_size=args.batch_per_gpu,
                        num_workers= args.workers,
                        pin_memory = True,
                        drop_last = True
                       )
    print(f"Data loaded")

    # === MODEL === #
    get_model = getattr(__import__("lib.arch.{}".format(args.arch), fromlist=["get_model"]), "get_model")
    setattr(args, "input_dim", dataset.num_node_features)
    setattr(args, "outpup_dim", dataset.num_classes)
    setattr(args, "activation_function", nn.functional.relu)
    model = get_model(args).cuda(args.gpu)
    model = nn.SyncBatchNorm.convert_sync_batchnorm(model) # use if model contains batchnorm.
    model = nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

    # === LOSS === #
    setattr(args, "loader", loader)
    from lib.core.loss import get_loss
    loss = get_loss(args).cuda(args.gpu)

    # === OPTIMIZER === #
    from lib.core.optimizer import get_optimizer
    optimizer = get_optimizer(model, args)

    # === TRAINING === #
    Trainer = getattr(__import__("lib.trainers.{}".format(args.trainer), fromlist=["Trainer"]), "Trainer")
    Trainer(args, loader, model, loss, optimizer).fit()

if __name__ == "__main__":
    main()


from src.Parameters import Parameters
from src.Roulette import Roulette
from src.EvalFunction import EvalFunction
from src.GrayConversion import GrayConversion
from src.Crossover import Xover
from src.Mutation import Mutation
import numpy as np
from  src.cProfiler import do_cprofile

class GeneticAlgorithm(Parameters):
    def __init__(self, steps_g, number_average, p_mutation):
        Parameters.__init__(self)
        self.g = 2  # number of generation
        self.steps_g = steps_g
        self.number_average = number_average
        self.p_mutation = p_mutation

    def begin_population(self):
        func = lambda x: ''.join(str(i) for i in np.array(np.random.randint(2, size=(self.chrom_length)), dtype=str))
        population = list(map(func, np.zeros(self.pop_size)))
        population = GrayConversion(population).chrom_to_gray_representation()
        return population

    def procedure(self):
        population = self.begin_population()
        tmp1 = np.zeros(self.steps_g)
        tmp2 = np.zeros(self.steps_g)
        for i, j in np.nditer([tmp1, tmp2], op_flags=['readwrite']):
            population = Roulette(population).choice()
            population = Xover(population).swap()
            population = Mutation(population, self.p_mutation).mutation()
            i[...], j[...] = EvalFunction(population).total_fit()
        return tmp1, tmp2

    @do_cprofile
    def average_results(self):
        tmp1, tmp2 = self.procedure()
        for i in range(self.number_average - 1):
            tmp = self.procedure()
            tmp1 += tmp[0]
            tmp2 += tmp[1]
        return np.transpose([tmp1 / self.number_average, tmp2 / self.number_average])

for i in range(10):
    np.savetxt('Gray_repr_xover:0.35_p_mut:' + str(0.03 + (i / 50)) + '.dat', GeneticAlgorithm(16000, 10, 0.04 + (i / 25)).average_results(), fmt=['%.9f', '%.9f'])

for i in range(5):
    np.savetxt('Gray_repr_xover:0.35_p_mut:' + str(0.4 + (i / 50)) + '.dat', GeneticAlgorithm(16000, 10, 0.04 + (i / 25)).average_results(), fmt=['%.9f', '%.9f'])


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

#  python-holidays
#  ---------------
#  A fast, efficient Python library for generating country, province and state
#  specific sets of holidays on the fly. It aims to make determining whether a
#  specific date is a holiday as fast and flexible as possible.
#
#  Author:  ryanss <ryanssdev@icloud.com> (c) 2014-2017
#           dr-prodigy <maurizio.montel@gmail.com> (c) 2017-2020
#  Website: https://github.com/dr-prodigy/python-holidays
#  License: MIT (see LICENSE file)

from datetime import date

from dateutil.easter import easter
from dateutil.relativedelta import relativedelta as rd, TH, FR, SA, SU

from holidays.constants import JAN, FEB, MAR, APR, MAY, SEP, OCT, \
    DEC
from holidays.holiday_base import HolidayBase

class Honduras(HolidayBase):
    # https://www.timeanddate.com/holidays/honduras/

    def __init__(self, **kwargs):
        self.country = "HND"
        HolidayBase.__init__(self, **kwargs)

    def _populate(self, year):
        # New Year's Day
        if self.observed and date(year, JAN, 1):
            self[date(year, JAN, 1)] = "Año Nuevo [New Year's Day]"

        # The Three Wise Men Day
        if self.observed and date(year, JAN, 6):
            name = "Día de los Reyes Magos [The Three Wise Men Day] (Observed)"
            self[date(year, JAN, 6)] = name

        # The Three Wise Men Day
        if self.observed and date(year, FEB, 3):
            name = "Día de la virgen de Suyapa [Our Lady of Suyapa] (Observed)"
            self[date(year, FEB, 3)] = name

        # The Father's Day
        if self.observed and date(year, MAR, 19):
            name = "Día del Padre [Father's Day] (Observed)"
            self[date(year, MAR, 19)] = name

        # Maundy Thursday
        self[easter(year) + rd(weekday=TH(-1))
             ] = "Jueves Santo [Maundy Thursday]"

        # Good Friday
        self[easter(year) + rd(weekday=FR(-1))
             ] = "Viernes Santo [Good Friday]"

        # Holy Saturday
        self[easter(year) + rd(weekday=SA(-1))
             ] = "Sábado de Gloria [Holy Saturday]"

        # Easter Sunday
        self[easter(year) + rd(weekday=SU(-1))
             ] = "Domingo de Resurrección [Easter Sunday]"

        # America Day
        if self.observed and date(year, APR, 14):
            self[date(year, APR, 14)] = "Día de las Américas [America Day]"

        # Labor Day
        if self.observed and date(year, MAY, 1):
            self[date(year, MAY, 1)] = "Día del Trabajo [Labour Day]"

        # Mother's Day
        may_first = date(int(year), 5, 1)
        weekday_seq = may_first.weekday()
        mom_day = (14 - weekday_seq)
        if self.observed and date(year, MAY, mom_day):
            str_day = "Día de la madre [Mother's Day] (Observed)"
            self[date(year, MAY, mom_day)] = str_day

        # Children's Day
        if self.observed and date(year, SEP, 10):
            name = "Día del niño [Children day] (Observed)"
            self[date(year, SEP, 10)] = name

        # Independence Day
        if self.observed and date(year, SEP, 15):
            name = "Día de la Independencia [Independence Day]"
            self[date(year, SEP, 15)] = name

        # Teacher's Day
        if self.observed and date(year, SEP, 17):
            name = "Día del Maestro [Teacher's day] (Observed)"
            self[date(year, SEP, 17)] = name

        # October Holidays are joined on 3 days starting at October 3 to 6.
        # Some companies work medium day and take the rest on saturday.
        # This holiday is variant and some companies work normally.
        # If start day is weekend is ignored.
        # The main objective of this is to increase the tourism.

        # https://www.hondurastips.hn/2017/09/20/de-donde-nace-el-feriado-morazanico/

        if year <= 2014:
            # Morazan's Day
            if self.observed and date(year, OCT, 3):
                self[date(year, OCT, 3)] = "Día de Morazán [Morazan's Day]"

            # Columbus Day
            if self.observed and date(year, OCT, 12):
                self[date(year, OCT, 12)] = "Día de la Raza [Columbus Day]"

            # Amy Day
            if self.observed and date(year, OCT, 21):
                str_day = "Día de las Fuerzas Armadas [Army Day]"
                self[date(year, OCT, 21)] = str_day
        else:
            # Morazan Weekend
            if self.observed and date(year, OCT, 3):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 3)] = name

            # Morazan Weekend
            if self.observed and date(year, OCT, 4):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 4)] = name

            # Morazan Weekend
            if self.observed and date(year, OCT, 5):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 5)] = name

        # Christmas
        self[date(year, DEC, 25)] = "Navidad [Christmas]"

class HN(Honduras):
    pass

class HND(Honduras):
    pass


import sys

sys.stdin = open("_암호생성기.txt")

for i in range(1,11):
    l = int(sys.stdin.readline())
    data = list(map(int,sys.stdin.readline().split()))
    cnt = 1
    temp = 0
    while True:
        temp = data[0]-cnt
        if cnt == 5:
            cnt = 0
        for j in range(len(data)-1):
            data[j] = data[j+1]
        data[len(data)-1] = temp
        cnt+=1
        if data[len(data)-1] <= 0:
            data[len(data)-1] = 0
            break
    print(f"#{i}", end=' ')
    for k in data:
        print(k,end=' ')
    print()

# deque

# from collections import deque

# for tc in range(10):
#     T = int(input())
#     queue = deque(map(int,input().split()))
#     cnt = 1
#     while queue[0] > 0:
#         check = queue.popleft()
#         check -= cnt
#         queue.append(check)
#         if check <= 0:
#             queue[-1] = 0
#             break
#         cnt = (cnt) % 5 + 1
#     queue = list(queue)
#     print(f"#{T}",*queue)