from typing import BinaryIO, Optional, Union

from typing_extensions import Protocol

from .batch_factory import Batch, BatchFactory
from .episode import EpisodeManager
from .io import dump_memory, load_memory
from .item import Item
from .step import StepBuffer
from .step_collector import StepCollector
from .transition_buffer import TransitionBuffer
from .transition_factory import TransitionFactory

class KioxProtocol(Protocol):
    def collect(
        self,
        observation: Item,
        action: Item,
        reward: Item,
        terminal: Union[bool, float],
        timeout: Optional[bool] = None,
    ) -> None:
        """Stores experience tuple.

        Args:
            observation: observation.
            action: action.
            reward: reward.
            terminal: terminal flag.
            timeout: timeout flag.

        """

    def get_step_buffer_size(self) -> int:
        """Returns number of stored tuples.

        Returns:
            number of tuples.

        """

    def get_transition_buffer_size(self) -> int:
        """Returns number of stored transitions.

        Returns:
            number of transitions.

        """

    def clip_episode(self) -> None:
        """Clips active episode.

        This method should be called whenever the current episode ends for
        timeout.

        """

    def sample(self, batch_size: int) -> Batch:
        """Samples transitions and returns mini-batch.

        Args:
            batch_size: batch size.

        Returns:
            mini-batch.

        """

    def copy_from(self, kiox: "KioxProtocol") -> None:
        """Copies steps and transitions from another Kiox object.

        Args:
            kiox: source Kiox object.

        """

    def save(self, f: BinaryIO) -> None:
        """Saves data as HDF5.

        Args:
            f: I/O-like object.

        """

    def load(self, f: BinaryIO) -> None:
        """Loads HDF5 data.

        Args:
            f: I/O-like object.

        """

class Kiox(KioxProtocol):
    """Kiox class.

    This class takes a single stream of experiences.

    .. code-block:: python

        kiox = Kiox(FIFOTransitionBuffer(1000), SimpleTransitionFactory())

        # collect data
        env = gym.make("CartPole-v0")
        obs = env.reset()
        while True:
            action = np.random.randint(2)
            next_obs, reward, terminal, _ = env.step(action)
            kiox.collect(obs, action, reward, terminal)
            if terminal:
                break
            obs = next_obs

        # sample mini-batch
        batch = kiox.sample(32)

        # save data
        with open("data.h5", "wb") as f:
            kiox.save(f)

        # load data
        kiox2 = Kiox(FIFOTransitionBuffer(1000), SimpleTransitionFactory())
        with open("data.h5", "rb") as f:
            kiox2.load(f)

    Args:
        transition_buffer: TransitionBuffer object.
        transition_factory: TransitionFactory object.
        n_steps: step size for multi-step learning. This corresponds to TD(N).
        gamma: discounted factor. If ``n_steps=1``, this value does not make
            any difference.

    """

    _step_buffer: StepBuffer
    _episode_manager: EpisodeManager
    _transition_buffer: TransitionBuffer
    _transition_factory: TransitionFactory
    _batch_factory: BatchFactory
    _step_collector: StepCollector

    def __init__(
        self,
        transition_buffer: TransitionBuffer,
        transition_factory: TransitionFactory,
        n_steps: int = 1,
        gamma: float = 0.99,
    ):
        self._step_buffer = StepBuffer()
        self._transition_buffer = transition_buffer
        self._transition_factory = transition_factory
        self._episode_manager = EpisodeManager(
            step_buffer=self._step_buffer,
            transition_buffer=self._transition_buffer,
        )
        self._batch_factory = BatchFactory(
            step_buffer=self._step_buffer,
            transition_buffer=transition_buffer,
        )
        self._step_collector = StepCollector(
            episode_manager=self._episode_manager,
            transition_factory=transition_factory,
            n_steps=n_steps,
            gamma=gamma,
        )

    def collect(
        self,
        observation: Item,
        action: Item,
        reward: Item,
        terminal: Union[bool, float],
        timeout: Optional[bool] = None,
    ) -> None:
        terminal = float(terminal) if isinstance(terminal, bool) else terminal
        self._step_collector.collect(
            observation=observation,
            action=action,
            reward=reward,
            terminal=terminal,
            timeout=timeout,
        )

    def get_step_buffer_size(self) -> int:
        return self._step_buffer.size()

    def get_transition_buffer_size(self) -> int:
        return self._transition_buffer.size()

    def clip_episode(self) -> None:
        self._step_collector.clip_episode()

    def sample(self, batch_size: int) -> Batch:
        return self._batch_factory.sample(batch_size)

    def copy_from(self, kiox: KioxProtocol) -> None:
        assert isinstance(kiox, Kiox)
        for episode in kiox.episode_manager.episodes:
            steps = episode.steps
            episode_length = len(steps)
            for i, step in enumerate(steps):
                self.collect(
                    observation=step.observation,
                    action=step.action,
                    reward=step.reward,
                    terminal=step.terminal,
                    timeout=(i + 1) == episode_length,
                )

    def save(self, f: BinaryIO) -> None:
        dump_memory(f, self._episode_manager.episodes)

    def load(self, f: BinaryIO) -> None:
        load_memory(f, self._step_collector)

    @property
    def episode_manager(self) -> EpisodeManager:
        return self._episode_manager

    @property
    def transition_buffer(self) -> TransitionBuffer:
        return self._transition_buffer

    @property
    def transition_factory(self) -> TransitionFactory:
        return self._transition_factory


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

# Font size and .tex background
params = {'text.usetex' : True,
          'font.size' : 15,
          'font.family' : 'lmodern'
          }
plt.rcParams.update(params)
fig = plt.figure()

# Find the experiment type
experiment_type = sys.argv[1]

truth = pd.read_csv('experiment-full/true-sin-wave.txt', delimiter=',').to_numpy()
param1 = pd.read_csv('experiment-full/sinwave-5.txt', delimiter=',').to_numpy()[:, 0]
param2 = pd.read_csv('experiment-full/sinwave-6.txt', delimiter=',').to_numpy()[:, 0]
param3 = pd.read_csv('experiment-full/sinwave-7.txt', delimiter=',').to_numpy()[:, 0]
param4 = pd.read_csv('experiment-full/sinwave-8.txt', delimiter=',').to_numpy()[:, 0]
param5 = pd.read_csv('experiment-full/sinwave-100.txt', delimiter=',').to_numpy()[:, 0]

data = np.array([param5, param1, param2, param3, param4])
#parameters = [ 0.001, 0.005, 0.05, 0.08 ]
parameters = [100, 500, 1000, 2000, 3000]
#parameters = [0.01, 0.02, 0.06, 0.1]
color_palette = ['royalblue', 'darkorange', 'crimson', 'limegreen', 'cyan']
xs = np.linspace(0, 900, len(truth))

for i in range(len(data)):
    plt.plot(xs, data[i], label='{} = {}'.format(experiment_type, parameters[i]), alpha=0.7, c=color_palette[i])

plt.plot(xs, truth[:, 0], label="Ground Truth", color='black', linewidth=2.0)

plt.ylabel('X-coordinate')
plt.xlabel('Time [s]')
plt.xlim(left=0)
plt.ylim(bottom=min(data.flatten()))
plt.title('True Movement vs. Predicted Movement Along X-Axis')
plt.legend(loc=2, prop={'size': 12})

# Save the file
fig.set_size_inches(7.2, 5)
plt.savefig('report-plots/svg/sinwave-{experiment}.svg'.format(experiment=experiment_type), dpi=1000, bbox_inches='tight')
plt.show()


from flask import url_for, flash, redirect, current_app, render_template
from main import db
from main.models import Appointment, Holiday, Post
from flask_login import current_user
from main.users.utils import send_email

def send_email_accept(id):
    appointment = Appointment.query.filter(Appointment.id == id).first()
    #SELECT * FROM Appointment WHERE id = {id} LIMIT 1;
    send_email('[Smile Plaza] Booked Appointment',
               sender=current_app.config['ADMINS'][0],
               recipients=[appointment.user_email],
               text_body=render_template('email/appointment_status_accepted.txt',
                                         appointment=appointment))

def send_email_reject(id):
    appointment = Appointment.query.filter(Appointment.id == id).first()
    #SELECT * FROM Appointment WHERE id = {id} LIMIT 1;
    send_email('[Smile Plaza] Booked Appointment',
               sender=current_app.config['ADMINS'][0],
               recipients=[appointment.user_email],
               text_body=render_template('email/appointment_status_rejected.txt',
                                         appointment=appointment))

def send_email_cancel(id):
    appointment = Appointment.query.filter(Appointment.id == id).first()
    #SELECT * FROM Appointment WHERE id = {id} LIMIT 1;
    send_email('[Smile Plaza] Booked Appointment',
               sender=current_app.config['ADMINS'][0],
               recipients=[appointment.user_email],
               text_body=render_template('email/appointment_status_cancelled.txt',
                                         appointment=appointment))

def send_email_received(id):
    appointment = Appointment.query.filter(Appointment.id == id).first()
    #SELECT * FROM Appointment WHERE id = {id} LIMIT 1;
    send_email('[Smile Plaza] Booked Appointment',
               sender=current_app.config['ADMINS'][0],
               recipients=[appointment.user_email],
               text_body=render_template('email/appointment_status_pending.txt',
                                         appointment=appointment))

def accept_action(appointmentID):
    appointment = Appointment.query.filter(Appointment.id == appointmentID).first()
    #SELECT * FROM Appointment WHERE id = {appointmentID} LIMIT 1;
    appointment.action = "ACCEPTED"
    #UPDATE Appointment SET action='ACCEPTED' WHERE id={appointmentID};
    db.session.commit()

def reject_action(appointmentID):
    appointment = Appointment.query.filter(Appointment.id == appointmentID).first()
    #SELECT * FROM Appointment WHERE id = {appointmentID} LIMIT 1;
    appointment.action = "REJECTED"
    appointment.status = "CANCELLED"
    #UPDATE Appointment SET action='REJECTED', status='CANCELLED' WHERE id={appointmentID};
    db.session.commit()

def cancel_status(appointmentID):
    appointment = Appointment.query.filter(Appointment.id == appointmentID).first()
    #SELECT * FROM Appointment WHERE id = {appointmentID} LIMIT 1;
    appointment.status = "CANCELLED"
    #UPDATE Appointment SET action='REJECTED', status='CANCELLED' WHERE id={appointmentID};
    db.session.commit()

def finish_status(appointmentID):
    appointment = Appointment.query.filter(Appointment.id == appointmentID).first()
    #SELECT * FROM Appointment WHERE id = {appointmentID} LIMIT 1;
    appointment.status = "FINISHED"
    #UPDATE Appointment SET status = 'FINISHED' WHERE id = {appointmentID};

    db.session.commit()

def holiday_status(selected_date_utc):
    if current_user.is_authenticated:
        holiday = Holiday(date=selected_date_utc, isHoliday="HOLIDAY")
        db.session.add(holiday)
        #INSERT INTO Holiday (date) VALUES ('{selected_date_utc}');
        db.session.commit()

        post = Post(title="No Clinic", content=f"No clinic for {selected_date_utc}, all appointments are cancelled.",
                    author=current_user)
        db.session.add(post)
        #INSERT INTO Post (title, content, author_id) VALUES ('No Clinic', CONCAT('No clinic for ' , '{selected_date_utc}', ', all appointments are cancelled.'), {current_user_id});
        db.session.commit()

        appointments = Appointment.query.filter(Appointment.date == selected_date_utc).all()
        #SELECT * FROM Appointment WHERE date = '{selected_date_utc}';
        print(f"Appointment: {appointments}")
        for appointment in appointments:
            appointment.action = "REJECTED"
            appointment.status = "CANCELLED"
            #UPDATE Appointment SET action = 'REJECTED', status = 'CANCELLED' WHERE date = '{selected_date_utc}';
            db.session.commit()
            send_email_cancel(appointment.id)

def add_appointment_to_database(selected_date_utc, selected_time, selected_service):
    appointment = Appointment(user_id=current_user.id, user_name=current_user.username,
                              user_email=current_user.email, user_contact = current_user.contact,
                              service = selected_service, date=selected_date_utc, time=selected_time)
    db.session.add(appointment)
    '''
    INSERT INTO Appointment (user_id, user_name, user_email, user_contact, service, date, time)
    VALUES ('{current_user_id}', '{current_user_username}', '{current_user_email}', '{current_user_contact}',
    '{selected_service}', '{selected_date_utc}', '{selected_time}');
    '''
    db.session.commit()
    flash('Your appointment has been added.', 'success')
    return redirect(url_for('main.customer_announcement'))


# Author: Olga Sotiriadi
# Lab:8
# Date:03.04.2022
# Purpose: Time Calculator
import pathlib
import tkinter as tk
import tkinter.ttk as ttk
import pygubu
import tkinter.messagebox as mb

PROJECT_PATH = pathlib.Path(__file__).parent
PROJECT_UI = PROJECT_PATH / "TimeCalc.ui"

class TimecalcApp:
    SEC_PER_DAY = 86400
    SEC_PER_HOUR = 3600
    SEC_PER_MIN = 60

    def __init__(self, master=None):
        # Boilerplate to build the tkinter interface based on TimeCalc.ui
        self.__builder = builder = pygubu.Builder()
        builder.add_resource_path(PROJECT_PATH)
        builder.add_from_file(PROJECT_UI)
        self.__mainwindow = builder.get_object('top_level', master)
        builder.connect_callbacks(self)

        # Save these UI elements as properties so that we can access them when
        # the user clicks on the calculate button.
        self.__days_entry = builder.get_object('days_entry', master)
        self.__hours_entry = builder.get_object('hours_entry', master)
        self.__minutes_entry = builder.get_object('minutes_entry', master)
        self.__seconds_entry = builder.get_object('seconds_entry', master)
        self.__result_entry_variable = builder.get_variable('result_entry_variable')

    def calculate(self):
        # Convert days, hours, minutes and seconds to seconds. If there's an error, display an error message.
        try:
            days = float(self.__days_entry.get())
            hours = float(self.__hours_entry.get())
            minutes = float(self.__minutes_entry.get())
            seconds = float(self.__seconds_entry.get())
            total_seconds = self.SEC_PER_DAY * days + self.SEC_PER_HOUR * hours + self.SEC_PER_MIN * minutes + seconds
            self.__result_entry_variable.set("{:.2f} seconds".format(total_seconds))
        except ValueError:
            mb.showerror(title="Error Calculating!", message="Please try again.")

    def run(self):
        self.__mainwindow.mainloop()

    def get_top_frame(self):
        # Return the top frame for the app so that it can be displayed in a tabbed notebook.
        return self.__mainwindow

    def clear(self):
        self.__days_entry.delete(0, tk.END)
        self.__hours_entry.delete(0, tk.END)
        self.__minutes_entry.delete(0, tk.END)
        self.__seconds_entry.delete(0, tk.END)

if __name__ == '__main__':
    root = tk.Tk()
    root.title("Time Calculator")
    app = TimecalcApp(root)
    app.run()


#!/usr/bin/env python3

import gym
import time
import procgen
import cv2
import numpy as np
import torch as th
from abc import abstractmethod
from typing import Tuple, Optional, Dict, Any
from pathlib import Path

_AGENT_COLOR: Tuple[int, int, int] = (127, 255, 127)
_WALL_COLOR: Tuple[int, int, int] = (191, 127, 63)
_WINDOW_SIZE: Tuple[int, int] = (7, 7)

class FlattenWrapper(gym.ObservationWrapper):
    def __init__(self, env: gym.Env):
        super().__init__(env)
        self.observation_space = gym.spaces.Box(
            low=0, high=255, shape=(
                np.prod(_WINDOW_SIZE) * 3,), dtype=np.uint8)

    def observation(self, observation: np.ndarray):
        # return observation.ravel()
        return (observation.ravel() / 255.0) - 0.5

class CustomEnvWrapper(gym.Wrapper):
    def __init__(self, env: gym.Env):
        super().__init__(env)
        self.env = env
        self.action_space = gym.spaces.Discrete(5)
        self.observation_space = gym.spaces.Box(
            low=0, high=255, shape=(_WINDOW_SIZE + (3,)), dtype=np.uint8)
        # left, down, up, right, place-marker
        self._actions = [1, 3, 5, 7, -1]
        self._markers = []
        self._history = set()

    def _get_pos(self, obs: np.ndarray) -> Tuple[int, int]:
        return tuple(np.argwhere(np.equal(obs, _AGENT_COLOR).all(
            axis=-1)).mean(axis=0).astype(np.int32))

    def _add_markers(self, obs: np.ndarray) -> np.ndarray:
        if len(self._markers) <= 0:
            return obs
        for m in self._markers:
            obs[m[0], m[1]] = np.bitwise_or(obs[m[0], m[1]], (255, 0, 0))
        return obs

    def _down_obs(self, obs: np.ndarray):
        # NOTE(ycho): [15/64] only works for `easy`
        return cv2.resize(obs, None, fx=(15 / 64), fy=(15 / 64),
                          interpolation=cv2.INTER_NEAREST_EXACT)

    def _crop_obs(self, obs: np.ndarray, loc: Tuple[int, int]):
        # cv2.imshow('obs_', obs)
        # cv2.waitKey(1)
        # return obs
        radius: Tuple[int, int] = (_WINDOW_SIZE[0] // 2, _WINDOW_SIZE[1] // 2)
        out = np.full(_WINDOW_SIZE + (3,),
                      255,
                      dtype=obs.dtype)

        idx0 = (max(0, loc[0] - radius[0]), max(0, loc[1] - radius[1]))
        idx1 = (min(loc[0] + radius[0] + 1, obs.shape[0]),
                min(loc[1] + radius[1] + 1, obs.shape[1]))
        offset = np.subtract(idx0, (loc[0] - radius[0], loc[1] - radius[1]))
        roi = obs[idx0[0]:idx1[0], idx0[1]:idx1[1]]
        out[offset[0]:offset[0] + roi.shape[0],
            offset[1]:offset[1] + roi.shape[1], :] = roi
        return out

    def _recolor_obs(self, obs: np.ndarray):
        obs[(obs == _AGENT_COLOR).all(axis=-1)] = (0, 255, 0)
        obs[(obs == _WALL_COLOR).all(axis=-1)] = (0, 0, 255)
        return obs

    def reset(self):
        self._markers = []
        self._history = set()
        obs = self._down_obs(self.env.reset())
        loc = self._get_pos(obs)
        self._history.add(loc)
        return self._recolor_obs(self._crop_obs(obs, loc))

    def step(self, action: int):
        # Take action.
        if action == 4:
            # NOTE(ycho): 4 = "null action" in procgen.
            obs, rew, done, info = self.env.step(4)
        else:
            obs, rew, done, info = self.env.step(self._actions[action])

        # Process observations.
        info['obs0'] = obs
        obs = self._down_obs(obs)
        loc = self._get_pos(obs)

        # For marker-placement actions,
        # track marker placements and add
        # extra penalty for rewards.
        if action == 4:
            rew -= 0.01
            self._markers.append(loc)
        obs = self._add_markers(obs)

        # Penalize each timestep.
        rew -= 0.01

        # Especially penalize revisits.
        if loc in self._history:
            rew -= 0.05

        # Track history, format observation and
        # and return results.
        self._history.add(loc)
        return self._recolor_obs(self._crop_obs(
            obs, loc)), rew, done, info

class MazeEnv(gym.Env):
    def __init__(self, env_config: Dict[str, Any]):
        super().__init__()
        env = gym.make('procgen:procgen-maze-v0',
                       # render_mode=render_mode,
                       # rand_seed=i,
                       start_level=0,
                       use_backgrounds=False,
                       use_monochrome_assets=True,
                       distribution_mode='easy',
                       # *args,
                       **env_config)
        env = CustomEnvWrapper(env)
        env = FlattenWrapper(env)
        self.env = env

        self.observation_space = gym.spaces.Box(
            low=-0.5, high=0.5, shape=(7 * 7 * 3,),
            dtype=np.float32)
        self.action_space = gym.spaces.Discrete(5)

    def reset(self):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)

def make_env(*args, **kwds) -> gym.Env:
    env = gym.make('procgen:procgen-maze-v0',
                   # render_mode=render_mode,
                   # rand_seed=i,
                   start_level=0,
                   # center_agent=True,
                   use_backgrounds=False,
                   use_monochrome_assets=True,
                   distribution_mode='easy',
                   # distribution_mode='memory'
                   *args, **kwds)
    env = CustomEnvWrapper(env)
    env = FlattenWrapper(env)
    return env


"""
created  by  hzwlxy  at 2018/7/5 10:59
__author__: 西瓜哥
__QQ__ : 120235331
__Note__：
"""
from datetime import datetime, timedelta
from flask import (session, _request_ctx_stack, request, current_app)
from flask_login import LoginManager as  _loginManager

from flask_login.signals import user_accessed, session_protected, user_loaded_from_cookie
from flask_login.config import COOKIE_NAME, AUTH_HEADER_NAME, COOKIE_DURATION, COOKIE_SECURE, COOKIE_HTTPONLY
from flask_login.utils import decode_cookie, encode_cookie

from flask_login_multi import _get_user,_user_context_processor

SESSION_KEYS = set(['user_id', 'remember', 'id', 'fresh', 'next'])

class LoginManager(_loginManager):

    def init_app(self, app, add_context_processor=True):

        app.login_manager = self
        app.after_request(self._update_remember_cookie)

        self._login_disabled = app.config.get('LOGIN_DISABLED', False)

        if add_context_processor:
            app.context_processor(_user_context_processor)

    @property
    def _get_endpoint(self):
        blueprint = getattr(request, 'blueprint', None)
        if blueprint is None:
            return 'user'
        return blueprint

    def reload_user(self, user=None):
        ctx = _request_ctx_stack.top
        if user is None:
            user_id = session.get(self._get_endpoint + '_user_id')
            if user_id is None:
                setattr(ctx, self._get_endpoint, self.anonymous_user())
            else:
                if self.user_callback is None:
                    raise Exception(
                        "No user_loader has been installed for this "
                        "LoginManager. Refer to"
                        "https://flask-login.readthedocs.io/"
                        "en/latest/#how-it-works for more info.")
                # 根据endpoint来判断当前登录是用户还是管理员
                user = self.user_callback(user_id, self._get_endpoint)
                if user is None:
                    setattr(ctx, self._get_endpoint, self.anonymous_user())
                else:
                    setattr(ctx, self._get_endpoint, user)
        else:
            setattr(ctx, self._get_endpoint, user)

    def _load_user(self):
        user_accessed.send(current_app._get_current_object())
        config = current_app.config

        if config.get('SESSION_PROTECTION', self.session_protection):
            deleted = self._session_protection()
            if deleted:
                return self.reload_user()

        is_missing_user_id = (self._get_endpoint + '_user_id') not in session
        if is_missing_user_id:
            cookie_name = self._get_endpoint + '_' + config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
            header_name = config.get('AUTH_HEADER_NAME', AUTH_HEADER_NAME)
            has_cookie = (cookie_name in request.cookies and
                          session.get('remember') != 'clear')

            if has_cookie:
                return self._load_from_cookie(request.cookies[cookie_name])
            elif self.request_callback:
                return self._load_from_request(request)
            elif header_name in request.headers:
                return self._load_from_header(request.headers[header_name])

        return self.reload_user()

    def _get_session_protection(self):
        app = current_app._get_current_object()
        return app.config.get('SESSION_PROTECTION', self.session_protection)

    def _session_protection(self):
        sess = session._get_current_object()
        ident = self._session_identifier_generator()
        app = current_app._get_current_object()
        mode = app.config.get('SESSION_PROTECTION', self.session_protection)

        # 用了endpoint区别之后，不能直接 用sess来判断session有没有值
        # 用sess.get(self._get_endpoint + '_user_id',None) 来判断当前用户session标志存不存在
        if sess.get(self._get_endpoint + '_user_id', None) and ident != sess.get(self._get_endpoint + '_id', None):
            if mode == 'basic' or sess.permanent:
                sess[self._get_endpoint + '_fresh'] = False
                session_protected.send(app)
                return False
            elif mode == 'strong':

                for k in SESSION_KEYS:
                    sess.pop(self._get_endpoint + '_' + k, None)
                sess[self._get_endpoint + '_remember'] = 'clear'
                session_protected.send(app)
                return True
        return False

    def _load_from_cookie(self, cookie):
        data = decode_cookie(cookie).split('||')
        user_id = data[0]
        mode = current_app._get_current_object().config.get('SESSION_PROTECTION', self.session_protection)
        if user_id is not None:
            session[self._get_endpoint + '_user_id'] = user_id
            session[self._get_endpoint + '_fresh'] = False
            if mode == 'strong':
                session[self._get_endpoint + '_id'] = data[1]

        self.reload_user()
        user = getattr(_request_ctx_stack.top, self._get_endpoint, None)

        if user is not None:
            app = current_app._get_current_object()
            user_loaded_from_cookie.send(app, user=_get_user())

    def _update_remember_cookie(self, response):

        if (self._get_endpoint + '_remember') not in session and current_app.config.get('REMEMBER_COOKIE_REFRESH_EACH_REQUEST'):
            session[self._get_endpoint + '_remember'] = 'set'

        if (self._get_endpoint + '_remember') in session:
            operation = session.pop((self._get_endpoint + '_remember'), None)

            if operation == 'set' and (self._get_endpoint + '_user_id') in session:
                self._set_cookie(response)
            elif operation == 'clear':
                self._clear_cookie(response)

        return response

    def _set_cookie(self, response):

        config = current_app.config
        cookie_name = self._get_endpoint + '_' + config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
        domain = config.get('REMEMBER_COOKIE_DOMAIN')
        path = config.get('REMEMBER_COOKIE_PATH', '/')
        secure = config.get('REMEMBER_COOKIE_SECURE', COOKIE_SECURE)
        httponly = config.get('REMEMBER_COOKIE_HTTPONLY', COOKIE_HTTPONLY)

        if (self._get_endpoint + '_remember_seconds') in session:
            duration = timedelta(seconds=session[self._get_endpoint + '_remember_seconds'])
        else:
            duration = config.get('REMEMBER_COOKIE_DURATION', COOKIE_DURATION)

        # 用户id 和 浏览器加密信息 存入cookie
        temp = (str(session[self._get_endpoint + '_user_id']), str(session[self._get_endpoint + '_id']))
        data = encode_cookie('||'.join(temp))

        if isinstance(duration, int):
            duration = timedelta(seconds=duration)

        try:
            expires = datetime.utcnow() + duration
        except TypeError:
            raise Exception('REMEMBER_COOKIE_DURATION must be a ' +
                            'datetime.timedelta, instead got: {0}'.format(
                                duration))

        response.set_cookie(cookie_name,
                            value=data,
                            expires=expires,
                            domain=domain,
                            path=path,
                            secure=secure,
                            httponly=httponly)

    def _clear_cookie(self, response):
        config = current_app.config
        cookie_name = self._get_endpoint + '_' + config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
        domain = config.get('REMEMBER_COOKIE_DOMAIN')
        path = config.get('REMEMBER_COOKIE_PATH', '/')
        response.delete_cookie(cookie_name, domain=domain, path=path)


import csv
import io
import json
import logging
import os
import shutil
import subprocess
import tempfile
import zipfile
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from copy import deepcopy
from datetime import timedelta, datetime
from typing import List, Tuple

import click
import numpy as np
from tabulate import tabulate
from termcolor import colored

import training_speech
from training_speech import utils, ffmpeg, sox, exceptions, vad

CURRENT_DIR = os.path.dirname(__file__)

logger = logging.getLogger()

@click.group()
def cli():
    pass

@cli.command()
@click.argument('source_name')
@click.option('-y', '--yes', is_flag=True, default=False, help='override existing transcript if any')
@click.option('--add-to-git/--no-add-to-git', is_flag=True, default=True)
def build_transcript(source_name, yes, add_to_git):
    source = training_speech.get_source(source_name)
    path_to_epub = os.path.join(CURRENT_DIR, 'data/epubs/', source['ebook'])

    path_to_transcript = os.path.join(CURRENT_DIR, f'data/transcripts/{source_name}.txt')
    if yes is False and os.path.isfile(path_to_transcript):
        click.confirm(text=f'{path_to_transcript} already exists. Override ?', default=False, abort=True)

    with open(path_to_transcript, 'w') as f:
        f.writelines(utils.read_epub(path_to_epub, path_to_xhtmls=source.get('ebook_parts', ['part1.xhtml'])))

    if add_to_git:
        subprocess.call(f'git add {path_to_transcript}'.split(' '))
    click.echo(f'transcript {path_to_transcript} added to git')

audio_player = None

def cut_fragment_audio(fragment: dict, input_file: str, output_dir: str=utils.CACHE_DIR, salt: str=None, force=False):
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)

    fragment_hash = utils.get_fragment_hash(fragment, salt=salt)
    path_to_fragment_audio = os.path.join(output_dir, f'{fragment_hash}.wav')
    if force or not os.path.isfile(path_to_fragment_audio):
        ffmpeg.cut(input_file, path_to_fragment_audio, from_=fragment['begin'], to=fragment['end'])
    return path_to_fragment_audio

def cut_fragments_audio(fragments: List[dict], input_file: str, output_dir: str=utils.CACHE_DIR):
    # generate fragments
    with click.progressbar(length=len(fragments), show_eta=True, label='cut audio into fragments') as bar:
        def _cut(f: dict):
            p = cut_fragment_audio(f, input_file, output_dir)
            bar.update(1)
            return p

        with ThreadPoolExecutor() as executor:
            return executor.map(_cut, fragments)

@cli.command()
@click.argument('source_name')
@click.option('-r', '--restart', is_flag=True, default=False, help='restart validation from scratch')
@click.option('-s', '--speed', default=1.3, help='set audio speed')
@click.option('-ar', '--audio-rate', default=16000)
@click.option('-nc', '--no-cache', is_flag=True, default=None)
@click.option('-f', '--fast', is_flag=True, default=False)
@click.option('--start', type=int, default=0)
def check_alignment(source_name, restart, speed, audio_rate, no_cache, fast, start):
    import inquirer
    source = training_speech.get_source(source_name)
    path_to_alignment = os.path.join(CURRENT_DIR, f'data/alignments/{source_name}.json')
    path_to_transcript = os.path.join(CURRENT_DIR, f'data/transcripts/{source_name}.txt')
    path_to_mp3 = os.path.join(CURRENT_DIR, 'data/mp3', source['audio'])

    if no_cache and os.path.isdir(utils.CACHE_DIR):
        shutil.rmtree(utils.CACHE_DIR)
        os.mkdir(utils.CACHE_DIR)

    # generate wav if do not exists yet
    with open(path_to_mp3, 'rb') as f:
        file_hash = utils.hash_file(f)
    path_to_wav = os.path.join(utils.CACHE_DIR, f'{file_hash}.wav')
    if not os.path.exists(path_to_wav):
        ffmpeg.convert(
            from_=path_to_mp3,
            to=os.path.join(utils.CACHE_DIR, f'{file_hash}.wav'),
            rate=audio_rate,
            channels=1
        )

    # retrieve transcript
    with open(path_to_transcript) as f:
        transcript = [l.strip() for l in f.readlines()]
    transcript = [l for l in transcript if l]  # rm empty lines

    # detect silences through VAD
    silences = vad.list_silences(path_to_wav=path_to_wav, frame_duration=20)

    if not restart and os.path.isfile(path_to_alignment):
        with open(path_to_alignment) as f:
            existing_alignment = json.load(f)
    else:
        existing_alignment = []

    alignment = utils.build_alignment(
        transcript=transcript,
        path_to_audio=path_to_wav,
        existing_alignment=existing_alignment,
        silences=silences,
        generate_labels=True,
    )

    def _check_alignment(index: int, alignment: List[dict]):
        click.clear()
        fragment = alignment[index]
        prev_fragments = alignment[max(i - 1, 0):i]
        next_fragments = alignment[i + 1:i + 3]

        print(colored(
            f'\nplaying #{i + 1:03d}: @@ {timedelta(seconds=fragment["begin"])}  {timedelta(seconds=fragment["end"])} ({fragment["end"] - fragment["begin"]:0.3f}) @@',  # noqa
            'yellow',
            attrs=['bold']
        ))
        if prev_fragments:
            for prev_ in prev_fragments:
                print(colored(prev_['text'], 'grey'))
        print(colored(fragment['text'], 'magenta' if fragment.get('warn') else 'green', attrs=['bold']))
        if next_fragments:
            for next_ in next_fragments:
                print(colored(next_['text'], 'grey'))

        todo = set()
        pool = ThreadPoolExecutor()

        def play_audio():
            path_to_audio = cut_fragment_audio(fragment, path_to_wav)
            global audio_player

            with sox.play(path_to_audio, speed=speed) as player:
                audio_player = player

        def play_audio_slow():
            path_to_audio = cut_fragment_audio(fragment, path_to_wav)
            global audio_player

            with sox.play(path_to_audio) as player:
                audio_player = player

        todo.add(pool.submit(play_audio))

        def ask_right_transcript(current: List[str]):
            new_text = click.edit(text='\n'.join(current), require_save=False)
            return [
                l.strip()
                for l in new_text.strip().split('\n')
                if l.strip()
            ]

        def ask_what_next():
            if prev_fragments:
                silence_before, silence_between, silence_after = utils.transition_silences(
                    prev_fragments[-1],
                    fragment,
                    silences
                )
                can_cut_start_on_prev_silence = silence_before is not None
                can_cut_start_on_next_silence = silence_after is not None
            else:
                can_cut_start_on_prev_silence = can_cut_start_on_next_silence = False
            if next_fragments:
                silence_before, silence_between, silence_after = utils.transition_silences(
                    fragment,
                    next_fragments[0],
                    silences
                )
                can_cut_end_on_prev_silence = silence_before is not None
                can_cut_end_on_next_silence = silence_after is not None
            else:
                can_cut_end_on_prev_silence = can_cut_end_on_next_silence = False

            try:
                next_: str = inquirer.prompt([
                    inquirer.List(
                        'next',
                        message="\nWhat should I do ?",
                        choices=(
                                ['approve', 'repeat'] +
                                (['go_back'] if prev_fragments else []) +
                                ['edit'] +
                                (['wrong_start__cut_on_previous_silence'] if can_cut_start_on_prev_silence else []) +
                                (['wrong_start__cut_on_next_silence'] if can_cut_start_on_next_silence else []) +
                                (['wrong_end__cut_on_previous_silence'] if can_cut_end_on_prev_silence else []) +
                                (['wrong_end__cut_on_next_silence'] if can_cut_end_on_next_silence else []) +
                                (['enable'] if fragment.get('disabled') else ['disable']) +
                                ['toggle_fast_mode','quit']),
                    ),
                ])['next']
            except TypeError:
                raise exceptions.QuitException
            except Exception:
                next_ = 'quit'

            global audio_player
            try:
                audio_player.kill()
            except:
                pass

            if next_ == 'repeat':
                todo.add(pool.submit(play_audio_slow))
                todo.add(pool.submit(ask_what_next))
            elif next_ == 'toggle_fast_mode':
                raise exceptions.ToggleFastModeException
            elif next_ == 'go_back':
                prev_fragments[-1].pop('disabled', None)
                prev_fragments[-1].pop('approved', None)
                raise exceptions.GoBackException
            elif next_ == 'edit':
                new_transcript = ask_right_transcript([t['text'] for t in prev_fragments + [fragment] + next_fragments])
                raise exceptions.SplitException(
                    start=i-len(prev_fragments),
                    end=i+len(next_fragments),
                    new_transcript=new_transcript,
                )

            elif next_ == 'wrong_start__cut_on_previous_silence':
                prev_fragment = prev_fragments[-1]
                silence_before, _, _ = utils.transition_silences(prev_fragment, fragment, silences)
                fragment['begin'] = round(max(silence_before[1] - 0.35, silence_before[1]), 3)
                prev_fragment['end'] = round(min(silence_before[0] + 0.35, silence_before[1]), 3)
                cut_fragment_audio(fragment, input_file=path_to_wav)
                cut_fragment_audio(fragment, input_file=path_to_wav)
                todo.add(pool.submit(play_audio))
                todo.add(pool.submit(ask_what_next))
            elif next_ == 'wrong_start__cut_on_next_silence':
                prev_fragment = prev_fragments[-1]
                _, _, silence_after = utils.transition_silences(prev_fragment, fragment, silences)
                prev_fragment['end'] = round(min(silence_after[0] + 0.35, silence_after[1]), 3)
                fragment['begin'] = round(max(silence_after[1] - 0.35, silence_after[0]), 3)
                cut_fragment_audio(prev_fragment, input_file=path_to_wav)
                cut_fragment_audio(fragment, input_file=path_to_wav)
                todo.add(pool.submit(play_audio))
                todo.add(pool.submit(ask_what_next))
            elif next_ == 'wrong_end__cut_on_previous_silence':
                next_fragment = next_fragments[0]
                silence_before, _, _ = utils.transition_silences(fragment, next_fragment, silences)
                fragment['end'] = round(min(silence_before[0] + 0.35, silence_before[1]), 3)
                next_fragment['begin'] = round(max(silence_before[1] - 0.35, silence_before[1]), 3)
                cut_fragment_audio(fragment, input_file=path_to_wav)
                cut_fragment_audio(next_fragment, input_file=path_to_wav)
                todo.add(pool.submit(play_audio))
                todo.add(pool.submit(ask_what_next))
            elif next_ == 'wrong_end__cut_on_next_silence':
                next_fragment = next_fragments[0]
                _, _, silence_after = utils.transition_silences(fragment, next_fragment, silences)
                fragment['end'] = round(min(silence_after[0] + 0.35, silence_after[1]), 3)
                next_fragment['begin'] = round(max(silence_after[1] - 0.35, silence_after[0]), 3)
                cut_fragment_audio(fragment, input_file=path_to_wav)
                cut_fragment_audio(next_fragment, input_file=path_to_wav)
                todo.add(pool.submit(play_audio))
                todo.add(pool.submit(ask_what_next))
            elif next_ == 'approve':
                fragment['approved'] = True
            elif next_ == 'pass':
                fragment.pop('approved', None)
            elif next_ == 'disable':
                fragment['disabled'] = True
                fragment.pop('approved', None)
            elif next_ == 'enabled':
                fragment['approved'] = True
                fragment.pop('disabled', None)
            elif next_ == 'quit':
                raise exceptions.QuitException
            else:
                raise NotImplementedError

        todo.add(pool.submit(ask_what_next))

        while todo:
            for future in as_completed(t for t in todo):
                todo.remove(future)
                future.result()

        pool.shutdown(wait=True)

    cut_fragments_audio(alignment, input_file=path_to_wav)

    # iterate over successive fragments
    start = start - 1
    i = start
    done = False
    while i < len(alignment) and not done:
        fragment = alignment[i]

        if start != i and (fragment.get('approved') or fragment.get('disabled')):
            click.echo(f'skip fragment#{i} {fragment["text"]}')
            i += 1
            continue

        if fast and i != 0 and i < len(alignment) - 1 and not fragment.get('warn'):
            fragment.update(
                approved=True,
                approved_auto=True,
            )
            click.echo(f'approve fragment#{i} {fragment["text"]}')
            i += 1
            continue

        try:
            _check_alignment(index=i, alignment=alignment)
        except exceptions.ToggleFastModeException:
            fast = not fast
            continue
        except exceptions.GoBackException:
            i -= 1
            continue
        except exceptions.QuitException:
            try:
                audio_player.kill()
            except:
                pass
            exit(1)
        except exceptions.SplitException as e:
            fragment.pop('approved', None)
            fragment.pop('disabled', None)
            audio_start: float = alignment[e.start]['begin']
            audio_end: float = alignment[e.end]['end']

            with tempfile.NamedTemporaryFile(suffix='.wav') as file_:
                sox.trim(path_to_wav, file_.name, from_=audio_start, to=audio_end)

                sub_alignment = utils.build_alignment(
                    transcript=e.new_transcript,
                    path_to_audio=file_.name,
                    existing_alignment=[
                        dict(
                            text=f['text'],
                            begin=f['begin'] - audio_start,
                            end=f['end'] - audio_start,
                            approved=f.get('approved', False),
                            disabled=f.get('disabled', False),
                        )
                        for f in alignment[e.start:e.end+1]
                    ],
                    silences=[
                        [max(s_start - audio_start, 0.), s_end - audio_start]
                        for s_start, s_end in silences
                        if s_end > audio_start and s_start < audio_end
                    ],
                    generate_labels=False,
                    language=source['language'],
                )

            for nf in sub_alignment:
                nf["begin"] += audio_start
                nf["end"] += audio_start

            alignment = (
                    alignment[:e.start] +
                    sub_alignment +
                    alignment[e.end+1:]
            )
            cut_fragments_audio(alignment, input_file=path_to_wav)
            i -= e.start

        # save progress
        with open(path_to_alignment, 'w') as dest:
            to_save = deepcopy(alignment)
            for f in to_save:
                f.pop('warn', None)
            json.dump(
                obj=to_save,
                fp=dest,
                sort_keys=True,
                indent=2,
            )
        with open(path_to_transcript, 'w') as f:
            f.writelines('\n'.join(f['text'] for f in alignment) + '\n')

        i += 1

MAPPINGS = [
    ('s3://audiocorp/epubs/', os.path.join(CURRENT_DIR, 'data/epubs/'), 'ebook'),  # epubs
    ('s3://audiocorp/mp3/', os.path.join(CURRENT_DIR, 'data/mp3/'), 'audio'),  # mp3
    ('s3://audiocorp/releases/', os.path.join(CURRENT_DIR, 'data/releases/'), 'releases'),
]

@cli.command()
@click.option('-s', '--source_name', default=None)
def download(source_name):
    for s3, local, key in MAPPINGS:
        if key in {'releases'}:
            continue
        local = os.path.abspath(local)
        options = ''
        if source_name:
            source = training_speech.get_source(source_name, validate=False)
            options += f' --exclude * --include {source[key]}'

        sync_cmd = f'aws s3 sync {s3} {local}{options}'
        print(sync_cmd)
        subprocess.call(sync_cmd.strip().split(' '))

@cli.command()
@click.option('-s', '--source_name', default=None)
@click.option('-r', '--releases', is_flag=True, default=False, help='releases only')
def upload(source_name, releases):
    for s3, local, key in MAPPINGS:
        local = os.path.abspath(local)
        options = ' --exclude .gitkeep'
        if (source_name and key in {'releases'}) or (releases and key != 'releases'):
            continue
        if key == 'releases':
            options += ' --acl public-read'
        else:
            options += ' --exclude *.zip'
        if source_name:
            source = training_speech.get_source(source_name)
            options += f' --exclude * --include {source[key]}'

        sync_cmd = f'aws s3 sync {local} {s3}{options}'
        print(sync_cmd)
        subprocess.call(sync_cmd.strip().split(' '))

@cli.command()
@click.option('-f', '--full', is_flag=True, default=False, help='display additionnal info')
def stats(full):
    sources = training_speech.sources()
    sources_data = []
    total_dur = timedelta(seconds=0)
    total_available = timedelta(seconds=0)
    total_count = 0
    per_language_count = defaultdict(float)
    per_language_dur = defaultdict(timedelta)
    per_language_available = defaultdict(timedelta)
    for name, metadata in sources.items():
        info = training_speech.source_info(name)
        path_to_mp3 = os.path.join(CURRENT_DIR, 'data/mp3', metadata['audio'])
        if full and os.path.isfile(path_to_mp3):
            mp3_duration = timedelta(seconds=ffmpeg.audio_duration(path_to_mp3))
        else:
            mp3_duration = None
        sources_data.append([
            name,
            info['status'],
            f'{int(info["progress"] * 100)} %',
            info['approved_count'],
            utils.format_timedelta(info['approved_duration']),
            utils.format_timedelta(mp3_duration) if mp3_duration else '?',
            metadata['language'],
        ])
        total_dur += info['approved_duration']
        total_count += info['approved_count']
        per_language_count[metadata['language']] += info['approved_count']
        per_language_dur[metadata['language']] += info['approved_duration']
        if mp3_duration:
            per_language_available[metadata['language']] += mp3_duration
            total_available += mp3_duration

    sources_data.append([])
    sources_data.append([
        'TOTAL',
        ''
        '',
        '',
        total_count,
        utils.format_timedelta(total_dur),
        utils.format_timedelta(total_available),
    ])
    for language, count in per_language_count.items():
        sources_data.append([
            f'TOTAL {language}',
            ''
            '',
            '',
            count,
            utils.format_timedelta(per_language_dur[language]),
            utils.format_timedelta(per_language_available[language]),
            language,
        ])
    print('\n' + tabulate(
        sources_data,
        headers=['Source', 'Status', 'Progress', '# speeches', 'Speeches duration', 'mp3 duration', 'Language'],
        tablefmt='pipe',
    ))

@cli.command()
@click.option('-r', '--audio-rate', default=16000)
@click.option('-l', '--language', type=click.Choice(['fr_FR']), default=None)
def release(audio_rate, language):
    per_language_sources = defaultdict(list)
    per_language_speakers = defaultdict(set)
    for name, metadata in training_speech.sources().items():
        info = training_speech.source_info(name)
        if info['status'] in {'DONE', 'WIP'}:
            per_language_sources[metadata['language']].append((name, metadata, info))
            per_language_speakers[metadata['language']].add(metadata['speaker'])

    today_str = datetime.now().isoformat()[:10]
    releases_data = []

    for source_language, sources in per_language_sources.items():
        if language and source_language != language:
            continue

        release_name = f'{today_str}_{source_language}'
        path_to_release = os.path.join(CURRENT_DIR, 'data/releases', f'{release_name}.zip')
        print(f'start building {release_name}')
        # generate fragments
        p_label = f'convert mp3 files to mono 16bits {audio_rate}Hz wav'
        with click.progressbar(length=len(sources), show_eta=True, label=p_label) as bar:
            with ThreadPoolExecutor() as executor:
                def _process_source(source_data: Tuple[str, dict, dict]):
                    source_name, metadata, _ = source_data
                    try:
                        path_to_alignment = os.path.join(CURRENT_DIR, f'data/alignments/{source_name}.json')
                        path_to_mp3 = os.path.join(CURRENT_DIR, 'data/mp3', metadata['audio'])
                        # generate wav if do not exists yet
                        with open(path_to_mp3, 'rb') as f:
                            file_hash = utils.hash_file(f)
                        path_to_wav = os.path.join(utils.CACHE_DIR, f'{file_hash}.wav')
                        if not os.path.exists(path_to_wav):
                            ffmpeg.convert(from_=path_to_mp3, to=path_to_wav, rate=audio_rate, channels=1)

                        with open(path_to_alignment) as file_:
                            source_fragments = [
                                dict(name=f'{source_name}_{i + 1:04d}', source_file=path_to_wav, **f)
                                for i, f in enumerate(json.load(file_))
                                if f.get('approved')
                            ]
                        bar.update(1)
                        return source_fragments
                    except Exception as e:
                        print(f'cannot process source {source_name}. {e}')
                        raise e

                fragments = []
                for source_fragments in executor.map(_process_source, sources):
                    # skip empty speeches and longer than 15s
                    fragments += [f for f in source_fragments if 0.1 <= f['end'] - f['begin'] <= 15]

        # generate fragments
        p_label = f'cut audio fragments'
        with click.progressbar(length=len(fragments), show_eta=True, label=p_label) as bar:
            with ThreadPoolExecutor() as executor:
                def _cut(f: dict):
                    try:
                        p = cut_fragment_audio(f, f['source_file'], salt=f['name'])
                        bar.update(1)
                        return p
                    except Exception as e:
                        print(f'cannot extract fragment {fragment["name"]}. {e}')

                audio_fragments_pathes = executor.map(_cut, fragments)

        p_label = f'generate {release_name}.zip file'
        with click.progressbar(length=len(fragments) + 1, show_eta=True, label=p_label) as bar:
            with zipfile.ZipFile(path_to_release, 'w') as zip_file:
                # create CSV
                string_buffer = io.StringIO()
                writer = csv.DictWriter(string_buffer, delimiter=',', fieldnames=['path', 'duration', 'text'])
                writer.writeheader()

                for fragment, tmp_audio in zip(fragments, audio_fragments_pathes):
                    archive_audio_path = f'{fragment["name"]}.wav'
                    zip_file.write(tmp_audio, arcname=archive_audio_path, compress_type=zipfile.ZIP_DEFLATED)
                    writer.writerow(dict(
                        path=archive_audio_path,
                        duration=round(fragment['end'] - fragment['begin'], 3),
                        text=fragment['text']
                    ))
                    bar.update(1)
                    os.unlink(tmp_audio)

                zip_file.writestr('data.csv', string_buffer.getvalue())
            bar.update(1)

        releases_data.append([
            f'[{release_name}](https://s3.eu-west-3.amazonaws.com/audiocorp/releases/{release_name}.zip)',
            len(fragments),
            len(per_language_speakers[source_language]),
            utils.format_timedelta(timedelta(seconds=round(sum(round(f['end'] - f['begin'], 3) for f in fragments)))),
            source_language,
        ])

    print('\n' + tabulate(
        releases_data,
        headers=['Name', '# speeches', '# speakers', 'Total Duration', 'Language'],
        tablefmt='pipe',
    ))

@cli.command()
@click.argument('source_name')
@click.argument('from_id', type=int)
@click.argument('to_id', type=int)
@click.option('-ar', '--audio-rate', default=16000)
def make_test(source_name, from_id, to_id, audio_rate):
    source = training_speech.get_source(source_name)
    path_to_alignment = os.path.join(CURRENT_DIR, f'data/alignments/{source_name}.json')
    path_to_transcript = os.path.join(CURRENT_DIR, f'data/transcripts/{source_name}.txt')
    path_to_mp3 = os.path.join(CURRENT_DIR, 'data/mp3', source['audio'])

    # generate wav if do not exists yet
    with open(path_to_mp3, 'rb') as f:
        file_hash = utils.hash_file(f)
    path_to_wav = os.path.join(utils.CACHE_DIR, f'{file_hash}.wav')
    if not os.path.exists(path_to_wav):
        ffmpeg.convert(
            from_=path_to_mp3,
            to=os.path.join(utils.CACHE_DIR, f'{file_hash}.wav'),
            rate=audio_rate,
            channels=1
        )

    # retrieve transcript
    with open(path_to_transcript) as f:
        transcript = [l.strip() for l in f.readlines()]
    transcript = [l for l in transcript if l]  # rm empty lines

    if os.path.isfile(path_to_alignment):
        with open(path_to_alignment) as f:
            existing_alignment = json.load(f)
    else:
        existing_alignment = []

    # detect silences
    silences = vad.list_silences(path_to_wav=path_to_wav, frame_duration=20)

    alignment = utils.build_alignment(
        transcript=transcript,
        path_to_audio=path_to_wav,
        existing_alignment=existing_alignment,
        silences=silences,
        generate_labels=True,
    )
    remaining = alignment[from_id - 1:to_id]

    with tempfile.NamedTemporaryFile(suffix='.wav') as file_:
        ffmpeg.cut(path_to_wav, file_.name, from_=remaining[0]['begin'], to=remaining[-1]['end'])
        with open(file_.name, 'rb') as f:
            file_hash = utils.hash_file(f)[:8]
        path_to_sub_audio = f'tests/assets/{file_hash}.wav'
        shutil.copy(file_.name, os.path.join(CURRENT_DIR, path_to_sub_audio))

    new_alignment = utils.build_alignment(
        transcript=[f['text'] for f in remaining],
        path_to_audio=path_to_sub_audio,
        existing_alignment=[],
        silences=vad.list_silences(path_to_wav=path_to_sub_audio, frame_duration=20),
        generate_labels=True,
    )

    print(BUILD_ALIGNMENT_TEST_TEMPLATE.format(
        file_hash=file_hash,
        source_name=source_name,
        from_=timedelta(seconds=remaining[0]['begin']),
        to=timedelta(seconds=remaining[-1]['end']),
        transcript='\n        '.join(f"'{f['text']}'," for f in new_alignment),
        alignment='\n        '.join("dict(begin={begin}, end={end}, text='{text}'),".format(**f) for f in new_alignment),
        vad_mode=utils.DEFAULT_VAD_MODE,
        vad_frame_duration=utils.DEFAULT_VAD_FRAME_DURATION,
    ))

BUILD_ALIGNMENT_TEST_TEMPLATE = """
    # {source_name} @@ {from_} {to} @@
    ('{file_hash}.wav', {vad_mode}, {vad_frame_duration}, [
        {transcript}
    ], [], [
        {alignment}
    ]),
"""

@cli.command()
@click.argument('source_name')
def source_stats(source_name):
    source = training_speech.get_source(source_name)
    path_to_alignment = os.path.join(CURRENT_DIR, f'data/alignments/{source_name}.json')
    path_to_transcript = os.path.join(CURRENT_DIR, f'data/transcripts/{source_name}.txt')
    path_to_mp3 = os.path.join(CURRENT_DIR, 'data/mp3', source['audio'])

    # generate wav if do not exists yet
    with open(path_to_mp3, 'rb') as f:
        file_hash = utils.hash_file(f)
    path_to_wav = os.path.join(utils.CACHE_DIR, f'{file_hash}.wav')
    if not os.path.exists(path_to_wav):
        ffmpeg.convert(
            from_=path_to_mp3,
            to=os.path.join(utils.CACHE_DIR, f'{file_hash}.wav'),
            rate=16000,
            channels=1
        )

    # retrieve transcript
    with open(path_to_transcript) as f:
        transcript = [l.strip() for l in f.readlines()]
    transcript = [l for l in transcript if l]  # rm empty lines

    if os.path.isfile(path_to_alignment):
        with open(path_to_alignment) as f:
            existing_alignment = json.load(f)
    else:
        existing_alignment = []

    # detect silences
    silences = vad.list_silences(path_to_wav=path_to_wav, frame_duration=utils.DEFAULT_VAD_FRAME_DURATION)

    alignment = utils.build_alignment(
        transcript=transcript,
        path_to_audio=path_to_wav,
        existing_alignment=existing_alignment,
        silences=silences,
        generate_labels=True,
    )

    transitions_durations = []
    fragments_durations = []
    for prev_fragment, next_fragment in zip(alignment[:-1], alignment[1:]):
        silence_before, silence_between, silence_after = \
            utils.transition_silences(prev_fragment, next_fragment, silences)
        if silence_between:
            transitions_durations.append(silence_between[1] - silence_between[0])
        fragments_durations.append(next_fragment['end'] - next_fragment['begin'])
    t_mean = np.array(transitions_durations).mean()
    t_std = np.array(transitions_durations).std()
    f_mean = np.array(fragments_durations).mean()
    f_std = np.array(fragments_durations).std()
    print('\n' + tabulate(
        [
            [
                'transition dur (s)',
                len(transitions_durations),
                round(t_mean, 3),
                round(t_std, 3),
                f'{round(t_mean - t_std, 3)} - {round(t_mean + t_std, 3)}',
                f'{round(t_mean - 2 * t_std, 3)} - {round(t_mean + 2 * t_std, 3)}',
                f'{max(0, round(t_mean - 3 * t_std, 3))} - {round(t_mean + 3 * t_std, 3)}',
                min(transitions_durations),
                max(transitions_durations),
            ],
            [
                'fragment dur (s)',
                len(fragments_durations),
                round(f_mean, 3),
                round(f_std, 3),
                f'{round(f_mean - f_std, 3)} - {round(f_mean + f_std, 3)}',
                f'{round(f_mean - 2 * f_std, 3)} - {round(f_mean + 2 * f_std, 3)}',
                f'{max(0, round(f_mean - 3 * f_std, 3))} - {round(f_mean + 3 * f_std, 3)}',
                min(fragments_durations),
                max(fragments_durations),
            ],
        ],
        headers=['Metric', 'count', 'avg', 'std', '70%', '95%', '95%', 'min', 'max'],
        tablefmt='pipe',
    ))

if __name__ == '__main__':
    cli()


import os
import numpy as np
from ReadFiles import read_vtk_file
import matplotlib.pyplot as plt
import Mathf as Mathf
from scipy.interpolate import griddata

_year = 365*24*3600

class NoFileError(Exception):
    pass

class FlagsError(Exception):
    pass

def SingularPoints(_data, col):
    shape = _data.shape
    sp = []     # find singular points
    preV = _data[0, col]
    for i in range(1, shape[0]):
        thisV = _data[i, col]
        if abs(thisV) > 1e-10 and abs(thisV - preV) / abs(thisV) > 0.1:
            sp.append(i)
        preV = thisV
    print("sp: ", sp)
    odata = np.zeros((shape[0] - len(sp), shape[1]))
    j = 0   # eliminate singular points
    for i in range(shape[0]):
        if i not in sp:
            odata[j, :] = _data[i, :]
            j += 1
    return odata

def LaterBetter(_data, _col):
    shape = _data.shape
    _t = _data[:, _col]
    now = 0
    Prev = shape[0] - 1
    Next = shape[0] - 2
    data = np.zeros(shape)
    data[now, :] = _data[Prev, :]
    while Next >= 0:
        if _t[Next] < _t[Prev]:
            Prev = Next
            now = now + 1
            data[now, :] = _data[Prev, :]
        Next = Next - 1
    data = data[0:now+1, :]
    data = data[::-1, :]
    return data

def bound_points(coord0, noz, nnoz = 1):
    coord1 = np.zeros((coord0.shape[0]//noz, coord0.shape[1]))
    for row in range(coord0.shape[0]):
        if row % noz == nnoz-1:
            nxy = row//noz
            coord1[nxy,:] = coord0[row,:]
    return coord1

def bound_elements(coord0,nox,noz):
    coord1 = np.zeros(((nox-1)**2,3))
    n0 = np.zeros(4)
    for i in range(nox-1):
        for j in range(nox-1):
            n0[0] = i*noz*nox+j*noz
            n0[1] = (i+1)*noz*nox+j*noz
            n0[2] = i*noz*nox+(j+1)*noz
            n0[3] = (i+1)*noz*nox+(j+1)*noz
            n = i*(nox-1)+j
            for k in n0:
                coord1[n,:] = coord1[n,:]+coord0[int(k),:]
            coord1[n,:] = coord1[n,:]/4.0
    return coord1

def line_prepender(filename, line):
    with open(filename, 'r+') as f:
        content = f.read()
        f.seek(0, 0)
        f.write(line.rstrip('\r\n') + '\n' + content)

def adjustThetaPhi(coord, index):
    for i in range(coord.shape[0]):
        coord[i, index[0]] = np.pi/2.0 - coord[i, index[0]]
        if coord[i, index[1]] < 0.0:
            coord[i, index[1]] = coord[i, index[1]] + 2.0 * np.pi

def adjustThetaPhi1(coord, index):
    for i in range(coord.shape[0]):
        coord[i, index[0]] = np.pi/2.0 - coord[i, index[0]]
        if coord[i, index[1]] < -np.pi:
            coord[i, index[1]] += 2.0 * np.pi
        elif coord[i, index[1]] > np.pi:
            coord[i, index[1]] -= 2.0 * np.pi

def expandBound(_coord, _field, index, limit):
    pi = np.pi
    n = 0
    coord1 = np.zeros(_coord.shape)
    field1 = np.zeros(_field.shape)
    for i in range(_coord.shape[0]):
        idxPhi = index[1]
        lon = _coord[i, idxPhi]
        idxLat = index[0]
        lat = _coord[i, idxLat]
        if lon < -pi + limit:
            coord1[n, :] = _coord[i, :]     # longtitude
            coord1[n, idxPhi] = _coord[i, idxPhi] + 2.0 * pi
            field1[n] = _field[i]
            n = n + 1
        elif lon > pi - limit:
            coord1[n, :] = _coord[i, :]
            coord1[n, idxPhi] = _coord[i, idxPhi] - 2.0 * pi
            field1[n] = _field[i]
            n = n + 1
        if lat < -pi / 2.0 + limit:
            coord1[n, :] = _coord[i, :]     # latitude
            coord1[n, idxLat] = -pi - lat
            if lon > 0.0:
                coord1[n, idxPhi] = lon - pi
            else:
                coord1[n, idxPhi] = lon + pi
            field1[n] = _field[i]
            n = n + 1
        if lat > pi / 2.0 - limit:
            coord1[n, :] = _coord[i, :]
            coord1[n, idxLat] = pi - lat
            if lon > 0.0:
                coord1[n, idxPhi] = lon - pi
            else:
                coord1[n, idxPhi] = lon + pi
            field1[n] = _field[i]
            n = n + 1
    coord = np.concatenate((_coord, coord1[0:n, :]), axis=0)
    field = np.concatenate((_field, field1[0:n, :]), axis=0)
    return coord, field

def middle_points(_nA):
    pA = np.zeros(_nA.size - 1)
    for i in range(0, _nA.size -1):
        pA[i] = (_nA[i] + _nA[i + 1]) / 2.0
    return pA

def EqualArea(n):
    n0 = 2 * n
    pi = np.pi
    cosTh = np.linspace(-1.0, 1.0, 2 * n + 1)
    cosTh = middle_points(cosTh)
    th = np.arccos(cosTh)
    lat = pi / 2.0 - th
    nphi = np.floor(2.0 * pi * n)
    phi = np.linspace(-pi, pi, nphi + 1)
    lon = middle_points(phi)
    area = 2 * pi / (n * nphi)
    # fig, ax = plt.subplots()
    # ax.plot(range(0, n0), lat, 'b*')
    # ax.set_ylim((-pi / 2.0, pi / 2.0))
    # fig.savefig('../eps/equalArea.jpg')
    return lat, lon, area

class CombineSurfFile():
    def __init__(self, fname, _inD, **kwargs):
        self.fname = fname
        nprocx = int(_inD['nprocx'])
        nprocz = int(_inD['nprocz'])
        self.nprocz = nprocz
        nodex = int(_inD['nodex'])
        nodez = int(_inD['nodez'])
        self.nproc = 12 * nprocx * nprocx * nprocz
        self.nox = (nodex+1)//nprocx
        self.noz = (nodez+1)//nprocz
        try:
            nodez = kwargs['nodez']
        except KeyError:
            nodez = 1
        if nodez < 2:
            self.procz = 0
            self.nnoz = nodez
        else:
            self.procz = (nodez-2)//(self.noz-1)
            self.nnoz = 2 + (nodez-2) % (self.noz-1)
        try:
            self.surfArea = kwargs["surfArea"]
        except KeyError:
            self.getSArea = False
        else:
            self.getSArea = True

    def __call__(self, route, case_name, step, procz, flags, **kwargs):
        print("Combine surf file, step = %s" % step)
        try:
            _oroute = kwargs['oroute']
        except KeyError:
            _oroute = route
        try:
            _area = kwargs['area']
        except KeyError:
            _area = False
        try:
            _expandLon = kwargs['expandLon']
        except KeyError:
            _expandLon = False
        if flags == "layer":
            self.CallLayer(route, case_name, step, procz)
        elif flags == "surf":
            self.CallSurf(route, case_name, step, procz,
                          obl=kwargs["obl"], expandLon=_expandLon, oroute=_oroute,
                          area=_area)
        else:
            raise FlagsError("flags %s is wrong" % (flags))
        try:
            eaMesh = kwargs['eaMesh']   # output an equal area mesh
        except KeyError:
            eaMesh = False
        if eaMesh:
            self.EaGrid(_oroute, case_name, procz, n=50, step=step)

    def CallLayer(self, route, case_name, step, procz):
        nprocxy = self.nproc//self.nprocz
        out_filename = os.path.join(route, "%s.lgmt.%s.%d.%d" %
                                           (case_name, self.fname, self.nodez, step))
        mm0 = np.zeros(2)
        print(out_filename)
        if os.path.isfile(out_filename):
            os.remove(out_filename)
        for i in range(nprocxy):
            field0, coord0 = self.Read1(route, case_name, step, i, self.procz)
            coord1 = bound_points(coord0, self.noz)
            field1 = bound_points(field0, self.noz, self.nnoz)
            mm1 = self.MaxMin(field1, 0)
            mm0 = mm1
            coord1 = Mathf.cart2sph(coord1)
            adjustThetaPhi(coord1, [1, 2])
            self.Output(out_filename, field1, coord1)
        line_prepender(out_filename, "%s %s" % (mm0[0], mm0[1]))
        pass

    # for MF files from Dr. MingMingli's code
    def CallSurf(self, route, case_name, step, procz, **kwargs):
        try:
            oroute = kwargs['oroute']
        except KeyError:
            oroute = route
        nprocxy = self.nproc//self.nprocz
        out_filename = os.path.join(oroute, "%s.MF0.%d" % (case_name, step))
        try:
            expandLon = kwargs['expandLon']
        except KeyError:
            expandLon = False
        if kwargs['obl'] is False and os.path.isfile(out_filename):
            print("%s exists" % out_filename)
            return
        mm0 = np.zeros(2)
        if os.path.isfile(out_filename):
            os.remove(out_filename)
        #   append integration as header
        try:
            _area = kwargs['area']
        except KeyError:
            _area = False
        if _area is not False:
            _getInteg = True
            _integrate = 0.0
        else:
            _getInteg = False
        maxV = 0.0
        for i in range(nprocxy):
            MF0, coord0 = self.Read(route, case_name, step, i)  # read MF file
            coord1 = bound_elements(coord0, self.nox, self.noz)
            coord1 = Mathf.cart2sph(coord1)
            adjustThetaPhi1(coord1, [1, 2])
            if self.getSArea:
                lel = int((self.nox - 1)**2.0)  # devide area of element
                surfA = self.surfArea[i * lel: (i + 1) * lel]
                for j in range(lel):
                    MF0[j, :] /= surfA[j]
            pmaxV = np.max(MF0[:, 0])
            if pmaxV > maxV:
                maxV = pmaxV
            if _getInteg:
                _integrate1 = self.IntegrateField(MF0, surfA)
                _integrate = _integrate + _integrate1
            if expandLon:
                coord1, MF0 = expandBound(coord1, MF0, [1, 2], np.pi/40.0)
            if i == 0:
                self.Output(out_filename, MF0, coord1)
            else:
                self.Output(out_filename, MF0, coord1, append=True)
        _line = "%.4e" % maxV
        if _getInteg:
            for i in range(_integrate.size):
                _line += "  %.4e" % _integrate[i]
        line_prepender(out_filename, _line)
        pass

    # generate an equal area grid
    def EaGrid(self, route, case_name, procz, **kwargs):
        pi = np.pi
        surfArea = 1740**2.0
        try:
            step = kwargs['step']
        except KeyError:
            step = 0
        else:
            print("Equal Area Grid: %d" % step)
            infile = os.path.join(route, "%s.MF0.%d" % (case_name, step))
            ofile = os.path.join(route, "%s.MF0_equalA.%d" % (case_name, step))
        try:
            total = kwargs['total']
        except KeyError:
            total = False
        if total:
            print("Equal Area Grid: Total")
            infile = os.path.join(route, "%s.MF_Total" % case_name)
            ofile = os.path.join(route, "%s.MF0_equalA_Total" % case_name)
        try:
            _n = kwargs['n']
        except KeyError:
            _n = 100
        data = np.genfromtxt(infile, skip_header=1)
        ilat = data[:, 0]
        ilon = data[:, 1]
        ifield = data[:, 2: data.shape[1] + 1]
        lat, lon, area = EqualArea(_n)
        lon, lat = np.meshgrid(lon, lat)
        size = lat.size
        lat.resize((size, ))   # grid output
        lon.resize((size, ))
        latAp = np.array([pi / 2.0, -pi / 2.0])     # add points near the polar
        lonAp = np.array([0.0, 0.0])
        lat = np.concatenate((lat, latAp))
        lon = np.concatenate((lon, lonAp))
        size = lat.size
        field = np.zeros((size, ifield.shape[1]))
        integrate = np.zeros(ifield.shape[1])
        for i in range(ifield.shape[1]):
            gfield = ifield[:, i]
            ffield = griddata((ilon, ilat), gfield, (lon, lat), method='linear')
            field[:, i] = ffield[:, ]
        maxV = np.max(field[:, 0])  # max value and total value
        for i in range(field.shape[0]):
            integrate += field[i, :] * area
        integrate *= surfArea
        coord = np.zeros((size, 3))
        coord[:, 1] = lat
        coord[:, 2] = lon
        self.Output(ofile, field, coord)
        line = "%.4e %.4e" % (maxV, area)
        for i in range(integrate.size):
            line += " %.4e" % integrate[i]
        line_prepender(ofile, line)

    #   .MF_Total file
    def TotalAmount(self, _cls, stepTuple, timeArray, **kwargs):
        print("TotalAmount")
        try:
            obl = kwargs['obl']
        except KeyError:
            obl = False
        route = _cls.route
        caseName = _cls.caseName
        timeScale = _cls.timeScale
        tScaleYear = timeScale/_year
        oFile = os.path.join(route, "%s.MF_Total" % caseName)
        if obl is False and os.path.isfile(oFile):
            return
        inFile = os.path.join(route, "%s.MF0.%d" % (caseName, stepTuple[0]))
        if os.path.isfile(inFile) is False:
            raise NoFileError(inFile)
        # Read in coordinate
        inData = np.genfromtxt(inFile, skip_header=1)
        shape = inData.shape
        coord0 = np.zeros((shape[0], 3))
        MM = np.zeros((shape[0], shape[1] - 2))
        coord0[:, 1:3] = inData[:, 0: 2]
        tt0 = timeArray[0]
        # melt rate * time
        for step in stepTuple:
            tt1 = timeArray[step]
            tt = tt1 - tt0
            tt0 = tt1
            inFile = os.path.join(route, "%s.MF0.%d" % (caseName, step))
            inData = np.genfromtxt(inFile, skip_header=1)
            MF = inData[:, 2: shape[1] + 1]
            MM = MM + MF * tt * tScaleYear
        maxV = np.max(MM[:, 0])
        self.Output(oFile, MM, coord0)
        _line = "%.4e" % maxV
        # surfA = self.surfArea   # Integration
        # _integrate = self.IntegrateField(MM, surfA)
        # for i in range(_integrate.size):
        #    _line += " %.4e" % (_integrate[i])
        line_prepender(oFile, _line)

    #   Intergration by element
    def IntegrateField(self, _field, _weight):
        integrate = np.zeros(_field.shape[1])
        for i in range(_field.shape[0]):
            integrate = integrate + _field[i, :] * _weight[i]
        return integrate

    # Read in MF file, flag is unused
    def Read(self, route, case_name, step, procxy, flag="add"):
        proc = procxy*self.nprocz
        filename=os.path.join(route,"%s.MF.%d.%d"%(case_name,proc,step))
        MF0 = np.genfromtxt(filename)
        for j in range(1,self.nprocz):
            proc = procxy*self.nprocz+j
            filename = os.path.join(route,"%s.proc%d.0.vts"%(case_name,proc))
            ddict,coord = read_vtk_file(filename)
            filename=os.path.join(route,"%s.MF.%d.%d"%(case_name,proc,step))
            MF1 = np.genfromtxt(filename)
            MF0 = MF0 + MF1
        return MF0, coord

    def Read1(self,route,case_name,step,procxy,procz):
    #Read in MF file, flag is unused#
        proc = procxy*self.nprocz+procz
        filename=os.path.join(route,"%s.proc%d.%d.vts"%(case_name,proc,step))
        ddict, coord = read_vtk_file(filename)
        field = ddict[self.fname]
        return field, coord

    def OutputHeader(selt,filename,mm):
        with open(filename,'w') as f:
            f.write("%.4e %.4e\n"%(mm[0],mm[1]))

    # output to a MF0 file for gmt
    def Output(self, filename, MF0, coord, **kwargs):
        cshape = coord.shape
        Mshape = MF0.shape
        try:
            _append = kwargs['append']
        except KeyError:
            _append = False
        if _append is True:
            _type = 'a'
        else:
            _type = 'w'
        with open(filename, _type) as f:
            for i in range(Mshape[0]):
                f.write("%.5e" % coord[i, 1])
                for j in range(2, cshape[1]):
                    f.write(" %.5e" % coord[i, j])
                for j in range(Mshape[1]):
                    f.write(" %.5e" % MF0[i, j])
                f.write("\n")

    # check data range
    def CheckRange(self, filename, MF0, coord, **kwargs):
        pass

    def MaxMin(self, MF, column):
        mm = np.zeros(2)
        mm[0] = np.min(MF[:, column])
        mm[1] = np.max(MF[:, column])
        return mm

# surf_area file
class SurfArea():
    def __init__(self, _inD, **kwargs):
        nprocx = int(_inD['nprocx'])
        nprocz = int(_inD['nprocz'])
        self.nprocx = nprocx
        self.nprocz = nprocz
        nodex = int(_inD['nodex'])
        nodez = int(_inD['nodez'])
        self.nproc = 12 * nprocx * nprocx * nprocz
        self.nox = (nodex+1)//nprocx
        self.noz = (nodez+1)//nprocz

    # read surf_area file and return a numpy array
    def __call__(self, route, prefix, **kwargs):
        nprocx = self.nprocx
        nprocz = self.nprocz
        nproc = 12 * nprocx**2 * nprocz
        ifile = os.path.join(route, "%s.surf_area.%d" % (prefix, nprocz - 1))  # read in data
        # print("SurfArea: %s" % ifile)
        if os.access(ifile, os.R_OK) is False:
            raise NoFileError(ifile)
        data = np.genfromtxt(ifile)
        for i in range(2 * nprocz - 1, nproc, nprocz):
            ifile = os.path.join(route, "%s.surf_area.%d" % (prefix, i))
            # print("SurfArea: %s" % ifile)
            if os.access(ifile, os.R_OK) is False:
                raise NoFileError(ifile)
            data1 = np.genfromtxt(ifile)
            data = np.concatenate((data, data1))
        try:
            col = kwargs['col']
        except KeyError:
            col = 0
        if len(data.shape) == 1:
            data.resize(data.shape[0], 1)   # fix 1-column bug
        return data[:, col]


import copy
import random

import pygame.image

from Custum_Components.bullet_etalons import BulletEtalons
from Custum_Components.enemies.Zdenek import Zdenek
from Custum_Components.enemies.boss_budai import BossBudai
from Custum_Components.enemies.creeper import Creeper
from Custum_Components.enemies.klickova import Klickova
from Custum_Components.enemies.sykora import Sykora
from Custum_Components.weapon import Weapon
from game_object import GameObject
import numpy as np
from Custum_Components.box_collider import BoxCollider
from Custum_Components.enemy import Enemy
import Custum_Components.basic_enemy_weapon as enemy_weapons

def get_spawn_pos_basic_side():
    return np.array([random.choice([-100.0, 700.0]), random.uniform(200, 450)])

class EnemyEtalons:
    def __init__(self, speed, health_q):
        bullets = BulletEtalons().bullets

        jura_weapon = enemy_weapons.BasicWeapon(bullet=bullets["basic"], fire_rate=1.4 * speed)
        jura = GameObject(np.array([0.0, 0.0]),
                          components=[BoxCollider(), jura_weapon,
                                      Enemy(max_health=7 * health_q, speed=0.6 * speed)],
                          img=pygame.image.load('sprites/Jura.png'),
                          tag='Enemy')

        # Velich
        velich = GameObject(np.array([0.0, 0.0]),
                            components=[BoxCollider(), Enemy(max_health=1 * health_q, speed=speed)],
                            img=pygame.image.load('sprites/Velich.png'),
                            tag='Enemy')
        velich.get_component(Enemy).get_spawn_pos = lambda: np.array([-100, random.randrange(300, 550)])
        velich.rotation = [1, 0]

        velich2 = copy.deepcopy(velich)
        velich2.get_component(Enemy).get_spawn_pos = lambda: np.array([700, random.randrange(300, 550)])
        velich2.rotation = [-1, 0]

        lopocha = GameObject(np.array([0.0, 0.0]),
                            components=[BoxCollider(), Enemy(max_health=6 * health_q, speed=2 * speed)],
                            img=pygame.image.load('sprites/Lopocha.jpeg'),
                            tag='Enemy')

        # Hrabalova
        h_weapon = enemy_weapons.EnemyTripleWeapon(bullet=bullets["basic"], fire_rate=0.7 * speed)
        hrabalova = GameObject(np.array([0.0, 0.0]),
                               components=[BoxCollider(), Enemy(max_health=15 * health_q, speed=0.1 * speed), h_weapon],
                               img=pygame.image.load('sprites/Hrabalova.png'),
                               tag='Enemy')

        b_weapon1 = enemy_weapons.EnemyTripleWeapon(bullet=bullets["basic"], fire_rate=speed, offset=[60, 0])

        budai = GameObject(np.array([0.0, 0.0]),
                           components=[BoxCollider(),
                                       BossBudai(max_health=500 * health_q, speed=0.5 * speed, boss=True),
                                       b_weapon1],
                           img=pygame.image.load('sprites/budai.png'),
                           tag='Enemy')

        rasek_weapon = enemy_weapons.BasicWeapon(bullet=bullets["explosive"], fire_rate=0.4 * speed)
        rasek = GameObject(np.array([0.0, 0.0]),
                           components=[BoxCollider(), Enemy(max_health=7 * health_q, speed=0.5 * speed), rasek_weapon],
                           img=pygame.image.load('sprites/Rasek.png'),
                           tag='Enemy')

        herman_weapon = enemy_weapons.Weapon(bullet=bullets["zig_zag"], fire_rate=0.4 * speed)
        herman = GameObject(np.array([0.0, 0.0]),
                            components=[BoxCollider(), Enemy(max_health=7 * health_q, speed=0.5 * speed),
                                        herman_weapon],
                            img=pygame.image.load('sprites/Herman.png'),
                            tag='Enemy')

        jedlicka = GameObject(np.array([0.0, 0.0]),
                              components=[BoxCollider(), Creeper(max_health=15 * health_q, speed=1.0 * speed,
                                                                 get_spawn_pos=get_spawn_pos_basic_side)],
                              img=pygame.image.load('sprites/Jedlicka.png'),
                              tag='Enemy')

        klickova = GameObject(np.array([0.0, 0.0]),
                              components=[BoxCollider(),
                                          Klickova(max_health=400 * health_q, speed=1.0 * speed, boss=True)],
                              img=pygame.image.load('sprites/Jedlicka.png'),
                              tag='Enemy')
        zdenda_weapon = enemy_weapons.BasicWeapon(bullet=bullets["basic"], fire_rate=1.4 * speed)
        zdenda = GameObject(np.array([0.0, 0.0]),
                            components=[BoxCollider(), Zdenek(max_health=200 * health_q, speed=1.0 * speed, boss=True),
                                        zdenda_weapon],
                            img=pygame.image.load('sprites/Zdenek.png'),
                            tag='Enemy')

        sykora = GameObject(np.array([0.0, 0.0]),
                            components=[BoxCollider(), Sykora(max_health=300 * health_q, speed=1.0 * speed, boss=True)],
                            img=pygame.image.load('sprites/Sykora.png'),
                            tag='Enemy')

        vondra_weapon1 = Weapon(bullet=bullets["basic"], fire_rate=1.0 * speed, direction=[0.0, -1.0])
        vondra_weapon2 = Weapon(bullet=bullets["basic"], fire_rate=1.0 * speed, direction=[1.0, -1.0])
        vondra_weapon3 = Weapon(bullet=bullets["basic"], fire_rate=1.0 * speed, direction=[-1.0, -1.0])

        vondra_ = GameObject(np.array([0.0, 0.0]),
                             components=[BoxCollider(), Enemy(max_health=4 * health_q, speed=0.5 * speed),
                                         vondra_weapon1, vondra_weapon2, vondra_weapon3],
                             img=pygame.image.load('sprites/Vondra.png'),
                             tag='Enemy')

        self.enemies = {
            "Jura": jura,
            "Velich": velich,
            "Velich2": velich2,
            "Vondra": vondra_,
            "Hrabalova": hrabalova,
            "budai": budai,
            "Rasek": rasek,
            "Herman": herman,
            "Jedlicka": jedlicka,
            "Zdenda": zdenda,
            "Sykora": sykora,
            "Klickova": klickova,
            "Zelenka": None,
            "Studenkova": None,
            "Lopocha": lopocha
        }


from sql_statement.authentication import Authentication
from sql_statement.commitment import Commitment
from sql_statement.ledger import Ledger
from sql_statement.score import Score
from sql_statement.user_login import UserLogin
from utils.date_utils import DateUtils

def init_scene_and_district(obj,scene_value,district_value):
    """初始化场景 区"""
    obj.set_scene(scene_value)
    obj.set_district(district_value)

def init_begin_time(obj,bt:str):
    """初始化开始时间"""
    obj.set_begin_time(bt)

def init_end_time(obj,et:str):
    """初始化结束时间"""
    obj.set_end_time(et)

def init_month_and_year(obj,month:int,year:int):
    """初始化月份 季度 年"""
    obj.set_month(month)

    # 根据月份生成对应的季度范围
    quarter_begin,quarter_end = DateUtils.get_quarter(month)
    obj.set_quarter_begin(quarter_begin)
    obj.set_quarter_end(quarter_end)

    obj.set_year(year)

def init_obj():
    """对5个统计对象的sql初始化"""
    user_login = UserLogin()
    score = Score()
    ledger = Ledger()
    commitment = Commitment()
    authentication = Authentication()
    return user_login,score,ledger,commitment,authentication

def init(objs,scene:int,district:str,year:int,month:int,begin_time:str,end_time:str):
    """初始化查询语句的条件变量"""
    # 为所有对象初始化场景，区
    for item in objs:
        init_scene_and_district(item,scene,district)
        # 台账初始化月份，年
        if isinstance(item,Ledger):
            init_month_and_year(item,month,year)
        # 自评 开始和结束时间
        if isinstance(item,Score):
            init_begin_time(item,begin_time)
            init_end_time(item,end_time)
        # 用户登录 开始和结束时间
        if isinstance(item,UserLogin):
            init_begin_time(item,begin_time)
            init_end_time(item, end_time)
        #  用户承诺 结束时间
        if isinstance(item,Commitment):
            init_end_time(item,end_time)

if __name__ == '__main__':
    objs = init_obj()
    print(objs)
    print(objs[0])
    print(isinstance(objs[0],UserLogin))
    # init(objs,1,'静安区',2023,10,'2023-10-01 00:00:00','2023-10-31 23:59:59')


"""
Character LSTM implementation (matches https://arxiv.org/pdf/1805.01052.pdf)
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class CharacterLSTM(nn.Module):
    def __init__(self, num_embeddings, d_embedding, d_out, char_dropout=0.0, **kwargs):
        super().__init__()

        self.d_embedding = d_embedding
        self.d_out = d_out

        self.lstm = nn.LSTM(
            self.d_embedding, self.d_out // 2, num_layers=1, bidirectional=True
        )

        self.emb = nn.Embedding(num_embeddings, self.d_embedding, **kwargs)
        self.char_dropout = nn.Dropout(char_dropout)

    def forward(self, chars_packed, valid_token_mask):
        inp_embs = nn.utils.rnn.PackedSequence(
            self.char_dropout(self.emb(chars_packed.data)),
            batch_sizes=chars_packed.batch_sizes,
            sorted_indices=chars_packed.sorted_indices,
            unsorted_indices=chars_packed.unsorted_indices,
        )

        _, (lstm_out, _) = self.lstm(inp_embs)
        lstm_out = torch.cat([lstm_out[0], lstm_out[1]], -1)

        # Switch to a representation where there are dummy vectors for invalid
        # tokens generated by padding.
        res = lstm_out.new_zeros(
            (valid_token_mask.shape[0], valid_token_mask.shape[1], lstm_out.shape[-1])
        )
        res[valid_token_mask] = lstm_out
        return res

class RetokenizerForCharLSTM:
    # Assumes that these control characters are not present in treebank text
    CHAR_UNK = "\0"
    CHAR_ID_UNK = 0
    CHAR_START_SENTENCE = "\1"
    CHAR_START_WORD = "\2"
    CHAR_STOP_WORD = "\3"
    CHAR_STOP_SENTENCE = "\4"

    def __init__(self, char_vocab):
        self.char_vocab = char_vocab

    @classmethod
    def build_vocab(cls, sentences):
        char_set = set()
        for sentence in sentences:
            if isinstance(sentence, tuple):
                sentence = sentence[0]
            for word in sentence:
                char_set |= set(word)

        # If codepoints are small (e.g. Latin alphabet), index by codepoint
        # directly
        highest_codepoint = max(ord(char) for char in char_set)
        if highest_codepoint < 512:
            if highest_codepoint < 256:
                highest_codepoint = 256
            else:
                highest_codepoint = 512

            char_vocab = {}
            # This also takes care of constants like CHAR_UNK, etc.
            for codepoint in range(highest_codepoint):
                char_vocab[chr(codepoint)] = codepoint
            return char_vocab
        else:
            char_vocab = {}
            char_vocab[cls.CHAR_UNK] = 0
            char_vocab[cls.CHAR_START_SENTENCE] = 1
            char_vocab[cls.CHAR_START_WORD] = 2
            char_vocab[cls.CHAR_STOP_WORD] = 3
            char_vocab[cls.CHAR_STOP_SENTENCE] = 4
            for id_, char in enumerate(sorted(char_set), start=5):
                char_vocab[char] = id_
            return char_vocab

    def __call__(self, words, space_after="ignored", return_tensors=None):
        if return_tensors != "np":
            raise NotImplementedError("Only return_tensors='np' is supported.")

        res = {}

        # Sentence-level start/stop tokens are encoded as 3 pseudo-chars
        # Within each word, account for 2 start/stop characters
        max_word_len = max(3, max(len(word) for word in words)) + 2
        char_ids = np.zeros((len(words) + 2, max_word_len), dtype=int)
        word_lens = np.zeros(len(words) + 2, dtype=int)

        char_ids[0, :5] = [
            self.char_vocab[self.CHAR_START_WORD],
            self.char_vocab[self.CHAR_START_SENTENCE],
            self.char_vocab[self.CHAR_START_SENTENCE],
            self.char_vocab[self.CHAR_START_SENTENCE],
            self.char_vocab[self.CHAR_STOP_WORD],
        ]
        word_lens[0] = 5
        for i, word in enumerate(words, start=1):
            char_ids[i, 0] = self.char_vocab[self.CHAR_START_WORD]
            for j, char in enumerate(word, start=1):
                char_ids[i, j] = self.char_vocab.get(char, self.CHAR_ID_UNK)
            char_ids[i, j + 1] = self.char_vocab[self.CHAR_STOP_WORD]
            word_lens[i] = j + 2
        char_ids[i + 1, :5] = [
            self.char_vocab[self.CHAR_START_WORD],
            self.char_vocab[self.CHAR_STOP_SENTENCE],
            self.char_vocab[self.CHAR_STOP_SENTENCE],
            self.char_vocab[self.CHAR_STOP_SENTENCE],
            self.char_vocab[self.CHAR_STOP_WORD],
        ]
        word_lens[i + 1] = 5

        res["char_ids"] = char_ids
        res["word_lens"] = word_lens
        res["valid_token_mask"] = np.ones_like(word_lens, dtype=bool)

        return res

    def pad(self, examples, return_tensors=None):
        if return_tensors != "pt":
            raise NotImplementedError("Only return_tensors='pt' is supported.")
        max_word_len = max(example["char_ids"].shape[-1] for example in examples)
        char_ids = torch.cat(
            [
                F.pad(
                    torch.tensor(example["char_ids"]),
                    (0, max_word_len - example["char_ids"].shape[-1]),
                )
                for example in examples
            ]
        )
        word_lens = torch.cat(
            [torch.tensor(example["word_lens"]) for example in examples]
        )
        valid_token_mask = nn.utils.rnn.pad_sequence(
            [torch.tensor(example["valid_token_mask"]) for example in examples],
            batch_first=True,
            padding_value=False,
        )

        char_ids = nn.utils.rnn.pack_padded_sequence(
            char_ids, word_lens, batch_first=True, enforce_sorted=False
        )
        return {
            "char_ids": char_ids,
            "valid_token_mask": valid_token_mask,
        }


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

whats = ("Q2", "Tk", "ang")
inpath = "qelma"

with PdfPages("qel_combo_ratio.pdf") as pdf:
    for what in whats:
        plt.figure()

        plt.title("numu, cc, E = 1000 MeV")

        x1, y1 = np.loadtxt("{}/numu1000_{}.txt".format(inpath, what), delimiter=' ', unpack=True)
        x2, y2 = np.loadtxt("{}/numu1200_{}.txt".format(inpath, what), delimiter=' ', unpack=True)
        x1r, y1r = np.loadtxt("{}/numu1000r_{}.txt".format(inpath, what), delimiter=' ', unpack=True)
        x2r, y2r = np.loadtxt("{}/numu1200r_{}.txt".format(inpath, what), delimiter=' ', unpack=True)

        plt.xlabel(what)
        plt.ylabel("#events")

        plt.plot(x1, (y1r - y1) / (y1r + y1), 'r-', label= 'A vs B to A')
        plt.plot(x2, (y2r - y2) / (y2r + y2), 'b-', label= 'B vs A to B')

        plt.legend()
        pdf.savefig()


"""
The following are implementations of algorithms that primarily operate on text. The list of implementations and their
purposes are

+--------------------------------------+-------------------------------------------------------------------+
| function                             | purpose                                                           |
+======================================+===================================================================+
| :func:`.anagrams`                    | Given an input word set, it returns a list of lists of words      |
|                                      | which are anagrams of each other.                                 |
+--------------------------------------+-------------------------------------------------------------------+

"""
import logging

# Set up the logger for the module
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s %(levelname)s %(name)s [%(lineno)d] %(message)s "
)

# Get the logger
logger: logging.Logger = logging.getLogger("algos.text")

def anagrams(word_set: set[str]) -> list[list[str]]:
    """
    Finds all anagrams of words contained within an input set of words.

    With an interactive Python session, you can run

    >>> from algos.text import anagrams
    >>> word_set = {'which', 'its', 'thing', 'drinking', 'can', 'car', 'that', 'cider', 'pain', 'and', 'by', 'below', \
'bowel', 'while', 'a', 'is', 'elbow', 'the', 'bending', 'in', 'during', 'an', 'arc', 'act', 'this', \
'cat', 'night', 'cried', 'rat', 'like', 'caused'}
    >>> anagrams(word_set)
    [['act', 'cat'], ['arc', 'car'], ['cider', 'cried'], ['bowel', 'elbow', 'below'], ['thing', 'night']]

    :raises TypeError: If the word_set is not a set.
    :raises TypeError: If the elements of the set are not all of type :class:`str` .
    :raises ValueError: If the word_set is empty.
    :param set[str] word_set: The set of words to find anagrams within.
    :return: A list of lists of anagrams.
    """
    # Create a dictionary that associates a word's signature to an array
    # containing all words with that signature.
    d: dict[str, list[str]] = {}

    # Declare the iterator.
    word: str

    # Raise TypeError if a set is not supplied.
    if not isinstance(word_set, set):
        logger.critical("anagrams - Incorrect Input Type")
        raise TypeError("Input Data Type Not Set")

    # Check that all element of set are of string type.
    if not all([isinstance(x, str) for x in word_set]):
        logger.critical("anagrams - Incorrect Elements Types")
        raise TypeError("Not All Elements of Type str")

    # Raise ValueError on empty set.
    if len(word_set) == 0:
        logger.critical("anagrams - Empty Input")
        raise ValueError("Empty Set")

    # For each word in the set of words.
    for word in word_set:
        # Calculate its signature. This is just each character sorted in alphabetical order as a single string.
        s: str = "".join(sorted(word))

        # If the signature has been encountered already
        if s in d:
            # Append the latest anagram to the list of anagrams with the same signature.
            d[s].append(word)
        # Otherwise this is the first occurrence of the signature.
        else:
            # Initialize the dictionary value with an array consisting of the word.
            d[s] = [word]

    # Return the lists of words for signatures that had more than one entry.
    return [d[s] for s in d if len(d[s]) > 1]