import inspect
import json
import os
from pathlib import Path
from typing import Dict, Literal
import warnings
warnings.warn('This submodule is deprecated. Please refer to `e15190.neutron_wall.adc_preprocessing` instead.', DeprecationWarning)

import duckdb as dk
import matplotlib as mpl
import matplotlib.pyplot as plt
import numdifftools as nd
import numpy as np
import pandas as pd
import ruptures as rpt
from sklearn.linear_model import RANSACRegressor
from sklearn.decomposition import PCA
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel
from scipy import interpolate

from e15190.runlog.query import Query
from e15190.neutron_wall.position_calibration import NWCalibrationReader
from e15190.neutron_wall.cache import RunCache
from e15190.utilities import fast_histogram as fh, misc, peak_finder, slicer, tables

class LightOutputCalibrator:
    DATABASE_DIR = '$DATABASE_DIR/neutron_wall/light_output_calibration'

    def __init__(self, AB: Literal['A', 'B']):
        self.AB = AB.upper()
        self.ab = self.AB.lower()
        self.df = None
        self.light_info = dict()

    @staticmethod
    def _get_daniele_root_files_dir(json_path=None) -> Path:
        if json_path is None:
            json_path = os.path.expandvars('$DATABASE_DIR/local_paths.json')
        with open(json_path) as file:
            return Path(json.load(file)['daniele_root_files_dir'])

    def read(self, runs, from_cache=True, **kwargs) -> pd.DataFrame:
        """Update ``self.df``"""
        rc = RunCache(
            src_path_fmt=str(self._get_daniele_root_files_dir() / r'CalibratedData_{run:04d}.root'),
            cache_path_fmt=str(Path(os.path.expandvars(self.DATABASE_DIR)) / r'cache/run-{run:04d}.db'),
        )

        kw = dict(
            from_cache=from_cache,
            sql_cmd='WHERE ' + ' AND '.join([
                'VW_multi == 0',
                f'NW{self.AB}_time_L - NW{self.AB}_time_R < 25',
                f'NW{self.AB}_time_L - NW{self.AB}_time_R > -25',
            ]),
            drop_columns=['VW_multi'],
            insert_run_index=True,
        )
        kw.update(kwargs)

        self.runs = [runs] if isinstance(runs, int) else runs
        self.df = rc.read(
            self.runs,
            {
                'VetoWall.fmulti'                           : 'VW_multi',
                f'NW{self.AB}.fnumbar'                      : f'NW{self.AB}_bar',
                f'NW{self.AB}.fLeft'                        : f'NW{self.AB}_total_L',
                f'NW{self.AB}.fRight'                       : f'NW{self.AB}_total_R',
                f'NW{self.AB}.ffastLeft'                    : f'NW{self.AB}_fast_L',
                f'NW{self.AB}.ffastRight'                   : f'NW{self.AB}_fast_R',
                f'NW{self.AB}.fTimeLeft'                    : f'NW{self.AB}_time_L',
                f'NW{self.AB}.fTimeRight'                   : f'NW{self.AB}_time_R',
            },
            **kw,
        )

        self.df = self.convert_64_to_32(self.df)
        self.df[f'NW{self.AB}_bar'] = self.df[f'NW{self.AB}_bar'].astype(np.int16)
        return self.df

    @staticmethod
    def convert_64_to_32(df):
        for col in df.columns:
            if df[col].dtype == np.float64:
                df[col] = df[col].astype(np.float32)
            elif df[col].dtype == np.int64:
                df[col] = df[col].astype(np.int32)
        return df

    def add_position(self, df=None, drop_time=True) -> pd.DataFrame:
        update_self = (df is None)
        if update_self:
            df = self.df

        bar = f'NW{self.AB}_bar'
        time_L = f'NW{self.AB}_time_L'
        time_R = f'NW{self.AB}_time_R'
        pos = f'NW{self.AB}_pos'

        calib_reader = NWCalibrationReader(self.AB)
        df_result = None
        for run, df_run in df.groupby('run'):
            pars = calib_reader(run)
            pars = pars.loc[df_run[bar]].to_numpy()
            df_run[pos] = pars[:, 0] + pars[:, 1] * (df_run[time_L] - df_run[time_R])
            df_result = pd.concat([df_result, df_run], axis=0)
        df_result = self.convert_64_to_32(df_result)

        if drop_time:
            df_result = df_result.drop([time_L, time_R], axis=1)
        if update_self:
            self.df = df_result
        return df_result

    @staticmethod
    def get_light_from_adc(adc, pos, a, b, c, d, e):
        result = 4.196 * adc / (a + b * pos + c * pos**2)
        result = d + result * e
        return result

    @staticmethod
    def get_adc_from_light(light, pos, a, b, c, d, e):
        result = (light - d) / (4.196 * e)
        result = result * (a + b * pos + c * pos**2)
        return result

    def read_light_calib_params(self, path=None) -> pd.DataFrame:
        if path is None:
            path = Path(os.path.expandvars(self.DATABASE_DIR)) / f'nw{self.ab}_pulse_height_calibration.dat'
        df = pd.read_csv(path, delim_whitespace=True, comment='#')
        return df.set_index('bar', drop=True)

    def _read_saturation_correction_params(self, path=None) -> pd.DataFrame:
        """Using Daniele's framework of Kuan's version."""
        if path is None:
            path = Path(os.path.expandvars(self.DATABASE_DIR)) / f'nw{self.ab}_saturation_correction.dat'
        df = pd.read_csv(path, delim_whitespace=True, comment='#')
        return df.set_index('bar', drop=True)

    @staticmethod
    def _randomize_columns(df, columns, seed=None):
        if isinstance(columns, str):
            columns = [columns]
        rand = np.random.RandomState(seed)
        df[columns] += rand.uniform(-0.5, 0.5, size=df[columns].shape)
        return df

    def randomize_ADC(self, seed=None):
        columns = [
            f'NW{self.AB}_total_L',
            f'NW{self.AB}_total_R',
            f'NW{self.AB}_fast_L',
            f'NW{self.AB}_fast_R',
        ]
        self.df = self._randomize_columns(self.df, columns, seed)
        self.df = self.convert_64_to_32(self.df)

    def analyze_log_of_light_ratio(self, light: Literal['total', 'fast'], bars=None, verbose=False):
        """
        """
        df = self.df.query(' & '.join([
            f'NW{self.AB}_{light}_L < 3500',
            f'NW{self.AB}_{light}_R < 3500',
        ]))[[
            f'NW{self.AB}_bar',
            f'NW{self.AB}_pos',
            f'NW{self.AB}_{light}_L',
            f'NW{self.AB}_{light}_R',
        ]]

        self.light_info.setdefault(light, dict())
        if bars is None:
            bars = range(1, 25)
        for bar in bars:
            if verbose:
                print(f'Analyzing bar-{bar:02d}...', flush=True)
            df_bar = dk.query(f'''
                SELECT * FROM df
                WHERE NW{self.AB}_bar == {bar}
            ''').df()
            self.light_info[light].setdefault(bar, dict())
            self.light_info[light][bar].update(self._analyze_log_of_light_ratio_per_bar(
                df_bar[f'NW{self.AB}_pos'],
                df_bar[f'NW{self.AB}_{light}_L'],
                df_bar[f'NW{self.AB}_{light}_R'],
            ))

    @staticmethod
    def _analyze_log_of_light_ratio_per_bar(pos, light_L, light_R):
        llr = LogOfLightRatio(pos, light_L, light_R)
        llr.fit_linear()
        llr.fit_wavy()
        return {
            'llr': llr,
            'attenuation_length': llr.attenuation_length,
            'gain_ratio': llr.gain_ratio,
            'log_light_ratio_spread': llr.log_light_ratio_spread,
        }

    @staticmethod
    def saturation_recovering(
        light_info: Dict[int, Dict[str, float]],
        bar_arr, pos, light_L, light_R,
        bars=None,
        threshold=4090,
    ):
        # populate scalars array
        scalars = np.zeros(len(bar_arr))
        if bars is None:
            bars = range(1, 25)
        for bar in bars:
            att_length = light_info[bar]['attenuation_length'][0]
            gain_ratio = light_info[bar]['gain_ratio'][0]
            scalars += (bar_arr == bar) * np.exp((2 / att_length) * pos + np.log(gain_ratio))

        # apply scalars conditionally to recover saturation
        recovered_light_L = np.where(
            (light_L > threshold) & (light_R < threshold) & (scalars > 1e-8),
            light_R / scalars,
            light_L,
        )
        recovered_light_R = np.where(
            (light_R > threshold) & (light_L < threshold) & (scalars > 1e-8),
            light_L * scalars,
            light_R,
        )
        return recovered_light_L, recovered_light_R

    @staticmethod
    def gain_matching(
        light_info: Dict[int, Dict[str, float]],
        bar_arr, light_L, light_R,
        bars=None,
    ):
        gain_ratios = np.zeros(len(bar_arr))
        if bars is None:
            bars = range(1, 25)
        for bar in bars:
            gain_ratio = light_info[bar]['gain_ratio'][0]
            gain_ratios += (bar_arr == bar) * gain_ratio
        gain_ratios += (gain_ratios < 1e-6) * 1.0

        matched_light_L = light_L * gain_ratios
        matched_light_R = light_R / gain_ratios
        return matched_light_L, matched_light_R

    def add_corrected_light(self, light: Literal['total', 'fast'], bars=None):
        light_L = self.df[f'NW{self.AB}_{light}_L']
        light_R = self.df[f'NW{self.AB}_{light}_R']

        light_L, light_R = self.saturation_recovering(
            self.light_info[light],
            self.df[f'NW{self.AB}_bar'],
            self.df[f'NW{self.AB}_pos'],
            light_L, light_R,
            bars,
        )

        light_L, light_R = self.gain_matching(
            self.light_info[light],
            self.df[f'NW{self.AB}_bar'],
            light_L, light_R,
            bars,
        )

        self.df[f'NW{self.AB}_{light}_L_corrected'] = light_L
        self.df[f'NW{self.AB}_{light}_R_corrected'] = light_R
        self.convert_64_to_32(self.df)

    def add_calibrated_light(self):
        """Calibrate geometric mean of corrected total ADCs into MeVee"""
        df_par = self.read_light_calib_params()
        df_pars = df_par.loc[self.df[f'NW{self.AB}_bar']]
        self.df[f'NW{self.AB}_light_GM'] = self.get_light_from_adc(
            np.sqrt(self.df[f'NW{self.AB}_total_L_corrected'] * self.df[f'NW{self.AB}_total_R_corrected']),
            np.array(self.df[f'NW{self.AB}_pos']),
            np.array(df_pars.a),
            np.array(df_pars.b),
            np.array(df_pars.c),
            np.array(df_pars.d),
            np.array(df_pars.e),
        )

class ParamIO:
    @classmethod
    def save_params(cls, calib):
        path = ParamIO.get_path(calib.runs, subdir='calib_params')
        path.parent.mkdir(parents=True, exist_ok=True)
        bar_info = dict()
        if path.is_file():
            bar_info = cls.df_to_dict(cls.read(path))
        for bar, info in calib.light_info['total'].items():
            bar_info[bar] = {
                'bar': bar,
                'att_length': info['attenuation_length'][0],
                'att_length_err': info['attenuation_length'][1],
                'gain_ratio': info['gain_ratio'][0],
                'gain_ratio_err': info['gain_ratio'][1],
                'log_light_ratio_spread': info['log_light_ratio_spread'],
            }
        bar_info = dict(sorted(bar_info.items(), key=lambda x: x[0]))
        df = pd.DataFrame(bar_info).T
        tables.to_fwf(
            df,
            path,
            comment=inspect.cleandoc(f'''
            # NW{calib.AB}
            # run(s): {calib.runs}
            # units:
            #     att_length: cm
            #     gain_ratio: unitless
            #     log_light_ratio_spread: unitless
            '''),
        )

    @staticmethod
    def read(path):
        return pd.read_csv(path, delim_whitespace=True, comment='#')

    @staticmethod
    def df_to_dict(df):
        df.set_index('bar', inplace=True, drop=True)
        return {bar: dict(bar=bar, **df.loc[bar]) for bar in df.index}

    @staticmethod
    def get_path(runs, subdir):
        path = Path(os.path.expandvars(LightOutputCalibrator.DATABASE_DIR)) / subdir
        if len(runs) == 1:
            path /= f'run-{runs[0]:04d}.dat'
        else:
            path /= f'{misc.runs_hash(runs)}.dat'
        return path

    @staticmethod
    def get_good_runs(min_run=None, max_run=None):
        min_run = min_run or 0
        max_run = max_run or 9999

        runs = np.array(list(range(min_run, max_run + 1)))
        return list(runs[Query.are_good(runs)])

    @classmethod
    def get_df_bars(cls, bars=None, runs=None, not_found_okay=True):
        bars = bars or range(1, 25)
        df_bars = {bar: [] for bar in bars}
        runs = runs or ParamIO.get_good_runs()
        for run in runs:
            path = cls.get_path([run], subdir='calib_params')
            if not path.is_file():
                if not not_found_okay:
                    raise FileNotFoundError(f'{path} does not exist')
                continue
            df_run = cls.read(path)
            df_run = df_run.set_index('bar', drop=True)
            for bar, holder in df_bars.items():
                holder.append([run, *df_run.loc[bar]])
        for bar, holder in df_bars.items():
            df_bars[bar] = pd.DataFrame(holder, columns=['run', *df_run.columns])
        return df_bars

    @staticmethod
    def median50_average(x):
        x = np.array(x)
        mid_range = np.quantile(x, [0.25, 0.75])
        mask = (x > mid_range[0]) & (x < mid_range[1])
        if sum(mask) == 0:
            return np.mean(x)
        return np.mean(x[mask])

    @staticmethod
    def median50_std(x):
        x = np.array(x)
        mid_range = np.quantile(x, [0.25, 0.75])
        mask = (x > mid_range[0]) & (x < mid_range[1])
        if sum(mask) < 4:
            return np.std(x)
        return np.std(x[mask])

    @classmethod
    def create_breakpoint_for_bar(cls, df_bar):
        data = np.array(df_bar[['att_length', 'gain_ratio', 'log_light_ratio_spread']].values)
        cpd = rpt.KernelCPD('rbf', min_size=1).fit(data)
        bp_indices = cpd.predict(pen=3)

        df_bp = []
        prev_i = 0
        for i_bp in bp_indices:
            rows = df_bar.iloc[prev_i:i_bp]
            att_length = cls.median50_average(rows['att_length'])
            gain_ratio = cls.median50_average(rows['gain_ratio'])
            llr_spread = cls.median50_average(rows['log_light_ratio_spread'])
            att_length_std = cls.median50_std(rows['att_length'])
            gain_ratio_std = cls.median50_std(rows['gain_ratio'])
            llr_spread_std = cls.median50_std(rows['log_light_ratio_spread'])
            df_bp.append([
                min(rows.run), max(rows.run),
                att_length, att_length_std,
                gain_ratio, gain_ratio_std,
                llr_spread, llr_spread_std,
            ])
            prev_i = i_bp
        return pd.DataFrame(
            df_bp,
            columns=[
                'run_start', 'run_stop',
                'att_length', 'att_length_std',
                'gain_ratio', 'gain_ratio_std',
                'log_light_ratio_spread', 'log_light_ratio_spread_std',
            ],
        )

    @classmethod
    def create_breakpoint_for_bars(cls, df_bars):
        result = None
        for bar, df_bar in df_bars.items():
            cuts = ['run < 3000', 'run > 3999']
            dfs_bp_bar = None
            for cut in cuts:
                df = df_bar.query(cut)
                if len(df) < 2: continue
                df_bp = cls.create_breakpoint_for_bar(df)
                df_bp.insert(0, 'bar', bar)
                dfs_bp_bar = pd.concat([dfs_bp_bar, df_bp])
            dfs_bp_bar = cls.fill_run_gaps(dfs_bp_bar)
            result = pd.concat([result, dfs_bp_bar])
        return result

    @staticmethod
    def fill_run_gaps(dfs_bp_bar):
        run_starts = dfs_bp_bar['run_start']
        run_stops = dfs_bp_bar['run_stop']
        run_ranges = np.vstack([run_starts, run_stops]).T
        for i in range(1, len(run_ranges)):
            if run_ranges[i, 0] != run_ranges[i - 1, 1] + 1:
                run_ranges[i, 0] = run_ranges[i - 1, 1] + 1

        # 3500 split
        i_4xxx = np.where(run_ranges[:, 0] > 3999)[0][0]
        if i_4xxx > 0:
            run_ranges[i_4xxx, 0] = 3501
            run_ranges[i_4xxx - 1, 1] = 3500

        result = dfs_bp_bar.copy()
        result['run_start'] = run_ranges[:, 0]
        result['run_stop'] = run_ranges[:, 1]
        return result

    @classmethod
    def create_and_save_breakpoints(cls, AB):
        df = cls.create_breakpoint_for_bars(cls.get_df_bars())
        path = Path(os.path.expandvars(LightOutputCalibrator.DATABASE_DIR)) / f'nw{AB.tolower()}_saturation_recovery.csv'
        df.to_csv(path, index=False)

class LogOfLightRatio:
    def __init__(self, pos, light_L, light_R):
        """Class to help analyzing the logarithm of light ratio.

        Data are assumed to be coming from one single NW bar. Experimental runs
        of the same nuclear reactions that are closeby in time could be chained
        together for more statistics.

        Both ``self.x`` and ``self.y`` will be defined as dimensionless
        quantities. The scaling factors are chosen empirically to make the
        ranges of x and y roughly the same.

        Parameters
        ----------
        light_L : array-like
            The light output of detected at the left side of the bar in ADC.
        light_R : array-like
            The light output of detected at the right side of the bar in ADC.
        pos : array-like
            The hit positions on the bar in centimeter.
        """
        self.x_scale = 100.0
        self.y_scale = 2.0
        mask = (light_L > 0) & (light_R > 0) & (pos > -150) & (pos < 150)
        self.df = pd.DataFrame({
            'x': pos[mask] / self.x_scale,
            'y': np.log(light_R[mask] / light_L[mask]) / self.y_scale,
        })

    @staticmethod
    def rotate(x, y, angle):
        return (
            x * np.cos(angle) + y * np.sin(angle),
            -x * np.sin(angle) + y * np.cos(angle),
        )

    def decorate_descale_xy(self, func):
        return lambda x: func(x / self.x_scale) * self.y_scale

    def fit_linear(self):
        """A robust linear fit on (x, y) using PCA and RANSAC.

        To improve robustness against outliers, PCA is first performed so that
        data points with 2nd components far from zero are removed. The remaining
        data are fed to RANSAC to find the most robust linear fit, updated as
        ``self.linear_fit``.
        """
        df = self.df[['x', 'y']].reset_index(drop=True)
        pca = PCA(n_components=2).fit(df)
        pca_xy = pca.transform(self.df[['x', 'y']])
        mask = (pca_xy[:, 1] > -0.4) & (pca_xy[:, 1] < 0.4)
        mask &= (self.df['x'] > -0.6) & (self.df['x'] < 0.6)
        ransac = RANSACRegressor(min_samples=0.5)
        ransac.fit(self.df[['x']][mask], self.df['y'][mask])
        self._linear_fit = np.polynomial.Polynomial([
            float(ransac.estimator_.intercept_),
            float(ransac.estimator_.coef_),
        ])
        self.linear_fit = self.decorate_descale_xy(self._linear_fit)

    def flatten(self, x, y):
        """Rotates (x, y) to zero slope of the linear fit.

        Rotation center is at origin, not the intersection between the linear
        fit and the y-axis. This function is only used intermediately for
        fitting the "wavy" relation.
        """
        angle = np.arctan(self._linear_fit.coef[1])
        return self.rotate(x, y, angle)

    def deflatten(self, x, y):
        """Rotates flattened (x, y) back to the original points.
        """
        angle = np.arctan(self._linear_fit.coef[1])
        return self.rotate(x, y, -angle)

    def fit_wavy_on_flat(self, return_xy=False):
        df = pd.DataFrame(
            np.vstack(self.flatten(self.df['x'], self.df['y'])).T,
            columns=['x_flat', 'y_flat'],
        )
        df = df.query('y_flat > -0.5 & y_flat < 0.5')

        xf, yf = [], []
        for x_range in slicer.create_ranges(-2.2, 2.2, width=0.1, step=0.02):
            arr = dk.query(f'''
                SELECT y_flat FROM df
                WHERE x_flat > {x_range[0]} AND x_flat < {x_range[1]}
            ''').df().y_flat
            if len(arr) < 30:
                continue
            finder = peak_finder.PeakFinderGaus1D(arr, hist_range=[-1, 1], hist_bins=200)
            xf.append(np.mean(x_range))
            yf.append(finder.get_highest_peak()[1])
        xf, yf = np.array(xf), np.array(yf)

        kernel = RBF(
            length_scale=1e-1,
            length_scale_bounds=(1e-3, 1e3),
        )
        kernel += WhiteKernel(
            noise_level=1e-5,
            noise_level_bounds=(1e-10, 1e-1),
        )
        self._gpr = GaussianProcessRegressor(kernel=kernel, n_restarts_optimizer=5)
        self._gpr.fit(xf[:, None], yf)
        self._wavy_fit_on_flat = lambda x: self._gpr.predict(x[:, None])
        return xf, yf if return_xy else None

    @staticmethod
    def extract_longest_increasing_subarray(x):
        n = len(x)
        i_longest, longest = 0, 0
        i_start, length = 0, 0
        for i in range(n - 1):
            if x[i] < x[i + 1]:
                length += 1
            else:
                if length > longest:
                    i_longest, longest = i_start, length
                i_start, length = i + 1, 0
        if length > longest:
            i_longest, longest = i_start, length
        return i_longest, i_longest + longest

    def fit_wavy(self):
        xf, _ = self.fit_wavy_on_flat(return_xy=True)
        xf = np.linspace(xf[0], xf[-1], 500)
        yf = self._wavy_fit_on_flat(xf)
        x, y = self.deflatten(xf, yf)

        # select longest monotonic x range
        i0, i1 = self.extract_longest_increasing_subarray(x)
        x = x[i0 : i1 + 1]
        y = y[i0 : i1 + 1]

        # extrapolate both ends linearly following the linear fit slope
        slope = self._linear_fit.coef[1]
        extrapolate = lambda x, m, x0, y0: m * (x - x0) + y0
        ext_L = lambda x_: extrapolate(x_, slope, x[0], y[0])
        ext_R = lambda x_: extrapolate(x_, slope, x[-1], y[-1])

        # interpolate the wavy curve
        x = np.hstack([x[0] - 1, x, x[-1] + 1])
        y = np.hstack([ext_L(x[0] - 1), y, ext_R(x[-1] + 1)])
        self._wavy_fit = interpolate.PchipInterpolator(x, y, extrapolate=True)
        self.wavy_fit = self.decorate_descale_xy(self._wavy_fit)

    def get_wavy_slopes(self, x):
        return nd.Derivative(self.wavy_fit, n=1)(x)

    @staticmethod
    def slope_to_attenuation_length(slope):
        return 2 / slope

    @staticmethod
    def intercept_to_gain_ratio(intercept):
        return np.exp(intercept)

    @property
    def attenuation_length_from_wavy(self):
        """In unit of centimeter"""
        x = np.linspace(-25, 25, 1000)
        slopes = self.get_wavy_slopes(x)
        att_lengths = self.slope_to_attenuation_length(slopes)
        return (
            np.mean(att_lengths),
            np.std(att_lengths),
        )

    @property
    def attenuation_length_from_linear(self):
        """In unit of centimeter"""
        slope = self.linear_fit(1) - self.linear_fit(0)
        return (
            self.slope_to_attenuation_length(slope),
            0.0
        )

    @property
    def attenuation_length(self):
        wavy = self.attenuation_length_from_wavy
        linear = self.attenuation_length_from_linear
        if np.allclose(wavy[0], linear[0], rtol=0.05):
            return wavy
        return (linear[0], wavy[1])

    @property
    def gain_ratio(self):
        """Gain right divided by gain left"""
        intercept = self.wavy_fit(0)
        xf, yf = self.flatten(0, intercept / self.y_scale)
        yf2, yf_err = self._gpr.predict([[xf]], return_std=True)
        if not np.isclose(yf, yf2, atol=1e-6):
            warnings.warn(f'GPR slope is not close to the linear slope at x = 0 cm')
        intercept_err = abs(intercept * yf_err / yf)

        gain_ratio = self.intercept_to_gain_ratio(intercept)
        gain_ratio_err = float(np.exp(gain_ratio) * intercept_err)

        return gain_ratio, gain_ratio_err

    @property
    def log_light_ratio_spread(self):
        y = self.df.query('x > -0.05 & x < 0.05').y
        return np.std(y) * self.y_scale

    def add_straightened_xy(self):
        dy = self._wavy_fit(self.df.x) - self._linear_fit(self.df.x)
        self.df['x_straight'] = self.df.x
        self.df['y_straight'] = self.df.y - dy

    def straighten(self, pos, log_light_ratio):
        return log_light_ratio - (self.wavy_fit(pos) - self.linear_fit(pos))

class LogOfLightRatioPlotter:
    def __init__(self, llr):
        self.llr = llr

    @property
    def x_plt(self):
        return np.linspace(-120, 120, 200)

    def _plot_hist2d(self, ax, df=None, straighten=False, **kwargs):
        if df is None:
            df = self.llr.df
        kw = dict(
            range=[[-150, 150], [-4, 4]],
            bins=[100, 100],
            cmap=plt.cm.viridis,
            norm=mpl.colors.LogNorm(vmin=1),
        )
        kw.update(kwargs)
        if straighten:
            x, y = 'x_straight', 'y_straight'
        else:
            x, y = 'x', 'y'
        fh.plot_histo2d(
            ax.hist2d,
            df[x] * self.llr.x_scale,
            df[y] * self.llr.y_scale,
            **kw,
        )
        return ax

    def _plot_linear(self, ax, **kwargs):
        kw = dict(
            color='purple',
            linewidth=1,
            linestyle='dashed',
        )
        kw.update(kwargs)
        ax.plot(self.x_plt, self.llr.linear_fit(self.x_plt), **kw)
        return ax

    def _plot_wavy(self, ax, **kwargs):
        kw = dict(
            color='red',
            linewidth=1,
        )
        kw.update(kwargs)
        ax.plot(self.x_plt, self.llr.wavy_fit(self.x_plt), **kw)
        return ax

    def plot(self, ax=None, df=None, straighten=False, wavy_fit=False, linear_fit=False):
        if ax is None:
            ax = plt.gca()
        self._plot_hist2d(ax, df=df, straighten=straighten)
        if wavy_fit:
            self._plot_wavy(ax)
        if linear_fit:
            self._plot_linear(ax)
        ax.set_xlabel('NW bar position (cm)')
        ax.set_ylabel(r'Log of light ratio $\ln(\mathcal{L}_R/\mathcal{L}_L)$')
        ax.set_xlim(-130, 130)
        ax.set_ylim(-3.5, 3.5)
        ax.axhline(0, color='gray', linestyle='dotted', linewidth=0.6)
        ax.axvline(0, color='gray', linestyle='dotted', linewidth=0.6)
        ax.annotate(
            '\n'.join([
            r'$\lambda = {:.1f} \pm {:.1f}$ cm'.format(*self.llr.attenuation_length),
            r'$g_R / g_L = {:.2f} \pm {:.2f}$'.format(*self.llr.gain_ratio),
            r'$\sigma_y = {:.2f}$'.format(self.llr.log_light_ratio_spread),
            ]),
            (0.05, 0.95),
            xycoords='axes fraction',
            ha='left',
            va='top',
            bbox=dict(boxstyle='round', facecolor='white', linewidth=0.3),
        )
        return ax

    @staticmethod
    def save_plots(calib, verbose=True):
        import matplotlib as mpl
        mpl.use('Agg')
        directory = Path(os.path.expandvars(calib.DATABASE_DIR)) / 'gallery'
        if len(calib.runs) == 1:
            run_title = f'run-{calib.runs[0]:04d}'
            directory /= f'{run_title}/'
        else:
            run_title = f'{misc.run_hash(calib.runs)}'
            directory /= f'{run_title}/'
            run_title = run_title.split('-h')[0]
        directory.mkdir(parents=True, exist_ok=True)

        for bar, info in calib.light_info['total'].items():
            if verbose:
                print(f'Plotting NW{calib.AB}-{bar:02d}...')
            path = directory / f'NW{calib.AB}-bar{bar:02d}.png'

            fig, ax = plt.subplots(figsize=(5, 4), facecolor='white', constrained_layout=True)
            llr = LogOfLightRatioPlotter(info['llr'])
            llr.plot(ax=ax, wavy_fit=True, linear_fit=True)
            ax.set_title(f'{run_title} NW{calib.AB}-{bar:02d}')
            fig.savefig(path, dpi=150, bbox_inches='tight')
            plt.draw()
            plt.close(fig)

class _Benchmark:
    """A development class to check if the code gives identical result to Daniele's framework."""
    def __init__(self):
        pass

    @staticmethod
    def read_daniele_root_file(path, tree_name='E15190'):
        import uproot
        branches = {
            'NWB.fnumbar': 'bar',
            'NWB.fLeft': 'total_L',
            'NWB.fRight': 'total_R',
            'NWB.fXcm': 'pos',
            'NWB.fGeoMeanSaturationCorrected': 'light_GM',
        }

        with uproot.open(str(path) + ':' + tree_name) as tree:
            df = tree.arrays(
                list(branches.keys()),
                entry_stop=10000,
                library='pd',
            )
        df.columns = list(branches.values())
        return df

    @staticmethod
    def get_saturation_corrected(df):
        """
        https://github.com/nscl-hira/E15190-Unified-Analysis-Framework/blob/837d3ad7a30a2678991bb5e6f9d7ddf0bc890473/NWPulseHeightCalibration.cpp#L223-L242
        """
        lo = LightOutputCalibrator('B')
        df_pul = lo.read_light_calib_params()
        df_sat = lo.read_saturation_correction_params()
        light_bm = []
        for _, entry in df.iterrows():
            bar = int(entry['bar'])
            p_sat = df_sat.loc[bar]
            if entry.total_L > 4090 and entry.total_R > 4090:
                light = np.sqrt(entry.total_L * entry.total_R)
            elif entry.total_L > 4090:
                light = np.exp(p_sat.p0 + p_sat.p1 * entry.pos + p_sat.p2 * entry.pos**2)
                light = np.sqrt(entry.total_R**2 / light)
            elif entry.total_R > 4090:
                light = np.exp(p_sat.p0 + p_sat.p1 * entry.pos + p_sat.p2 * entry.pos**2)
                light = np.sqrt(entry.total_L**2 * light)
            else:
                light = np.sqrt(entry.total_L * entry.total_R)

            p_pul = df_pul.loc[bar]
            light += np.random.uniform(-0.5, 0.5)
            light /= p_pul.a + p_pul.b * entry.pos + p_pul.c * entry.pos**2
            light = p_pul.d + 4.196 * p_pul.e * light

            light_bm.append(light)
        df['light_bm'] = light_bm
        return df

    @staticmethod
    def check_daniele_kuan(plot=False):
        from e15190 import DATABASE_DIR
        import json
        with open(DATABASE_DIR / 'local_paths.json') as file:
            directory = Path(json.load(file)['daniele_root_files_dir'])

        df = _Benchmark.read_daniele_root_file(directory / 'CalibratedData_4083.root')
        df = _Benchmark.get_saturation_corrected(df)

        diff = df.light_bm - df.light_GM
        diff = diff[diff < 100]

        if plot:
            fig, ax = plt.subplots()
            ax.scatter(df.light_GM, df.light_bm, s=0.5)
            ax.set_xlim(0, 100)
            ax.set_ylim(0, 100)
            plt.show()

        assert np.max(np.abs(diff)) < 2

class _MainUtilities:
    @staticmethod
    def get_args():
        parser = argparse.ArgumentParser(
            description='Calibrate the light output for neutron wall.',
            formatter_class=argparse.RawTextHelpFormatter,
        )
        parser.add_argument(
            'AB',
            type=str,
            choices=['A', 'B'],
            help='"A" or "B". This selects "NWA" or "NWB".',
        )
        parser.add_argument(
            'runs',
            nargs='+',
            help=inspect.cleandoc('''
                Runs to calibrate.

                Consecutive runs can be specified as a range. Here is an
                example:
                    > ./light_output_calibration.py B 8-10 11 20 2-5
                This will calibrate the runs 8, 9, 10, 11, 20, 2, 3, 4, 5, on
                NWB.
            '''),
        )
        parser.add_argument(
            '-b', '--bars',
            nargs='+',
            help=misc.MainUtilities.wrap('''
                The bar number(s) to calibrate. If not specified, all bars from
                bar 1 to 24 will be analyzed. Similar to "runs", dash can be
                used to specify ranges.
            '''),
            default=['1-24'],
        )
        parser.add_argument(
            '-c', '--no-cache',
            help=misc.MainUtilities.wrap('''
                When this option is given, the script will ignore the caches.
                All data will hence be read from the (Daniele's) ROOT files. New
                cache files will then be created, overwritting the old ones.
            '''),
            action='store_true',
        )
        parser.add_argument(
            '-d', '--debug',
            help=misc.MainUtilities.wrap('''
                When this option is given, no output would be saved. This
                includes both calibration parameters and gallery. The cache
                files, however, will still be updated.
            '''),
            action='store_true',
        )
        parser.add_argument(
            '-o', '--output',
            help=misc.MainUtilities.wrap('''
                The output directory for resultant ROOT files. If not given, the default is
                "$DATABASE_DIR/neutron_wall/light_output_calibration/".
            '''),
            default="$DATABASE_DIR/neutron_wall/light_output_calibration/",
        )
        parser.add_argument(
            '-s', '--silence',
            help='To silent all status messages.',
            action='store_true',
        )
        args = parser.parse_args()
        args.runs = misc.MainUtilities.parse_runs(
            args.runs,
            Query.are_good,
            verbose=(not args.silence),
        )
        args.bars = misc.MainUtilities.parse_bars(args.bars)
        return args

if __name__ == '__main__':
    import argparse
    args = _MainUtilities.get_args()

    LightOutputCalibrator.DATABASE_DIR = args.output
    calib = LightOutputCalibrator(args.AB)
    for run in args.runs:
        try:
            calib.read(
                run,
                from_cache=(not args.no_cache),
                verbose=(not args.silence),
            )
            calib.add_position()
            calib.randomize_ADC()
            calib.analyze_log_of_light_ratio(
                'total',
                bars=args.bars,
                verbose=(not args.silence),
            )
            if not args.debug:
                ParamIO.save_params(calib)
                LogOfLightRatioPlotter.save_plots(calib)
        except Exception as err:
            print(f'Error in run {run}:\n{err}')

{'uproot': 'uproot', 'DATABASE_DIR': 'e15190.DATABASE_DIR', 'json': 'json'}.read_daniele_root_file

{'uproot': 'uproot', 'DATABASE_DIR': 'e15190.DATABASE_DIR', 'json': 'json'}.get_saturation_corrected


'''
Created on 24.10.2015

@author: michael
'''
import unittest
from unittest.mock import MagicMock

from alexandriabase import baseinjectorkeys
from alexpresenters.MessageBroker import Message, CONF_EVENT_CHANGED, \
    REQ_SET_EVENT
from alexpresenters.Module import PresentersModule
from ddt import ddt, data, unpack
from integration.components.references.basereferenceintegrationtest \
    import BaseReferenceIntegrationTest
from tkgui import guiinjectorkeys
from tkgui.References import EventCrossReferencesView

@ddt
class EventCrossReferencesPresenterTest(BaseReferenceIntegrationTest):

    def setUp(self):
        super().setUp()
        self.injector = self.get_injector(PresentersModule())
        self.event_dao = self.injector.get(baseinjectorkeys.EVENT_DAO_KEY)
        self.presenter = self.injector.get(guiinjectorkeys.EVENT_CROSS_REFERENCES_PRESENTER_KEY)
        self.presenter.view = MagicMock(spec=EventCrossReferencesView)

    def testReceiveMessage(self):

        # Construct message
        event = self.event_dao.get_by_id(1940000001)
        message = Message(CONF_EVENT_CHANGED, event=event)

        # Message
        self.presenter.receive_message(message)

        # Assert
        self.assertTrue(self.presenter.view.current_event == event)
        self.assertTrue(len(self.presenter.view.items), 2)
        self.assertIn(self.event_dao.get_by_id(1950000001), self.presenter.view.items)
        self.assertIn(self.event_dao.get_by_id(1960013001), self.presenter.view.items)

    def testGotoEvent(self):
        event = self.event_dao.get_by_id(1940000001)
        self.presenter.view.selected_item = event
        self.presenter.goto_event()
        self.assertEqual(len(self.received_messages), 1)
        self.assertEqual(self.received_messages[0].key, REQ_SET_EVENT)
        self.assertEqual(self.received_messages[0].event, event)

    def testGotoEventNothingSelected(self):
        self.presenter.view.selected_item = None # If we wouldn't intialize, we'd get a mock
        self.presenter.goto_event()
        self.assertEqual(len(self.received_messages), 0)

    @data([1940000001, (1950000001, 1960013001, 1961050101)],
          [None, (1961050101,)])
    @unpack
    def testAddNewCrossReference(self, event_id, reference_ids):

        self.set_current_event(event_id)
        self.presenter.view.new_cross_reference_event = self.event_dao.get_by_id(1961050101)
        self.presenter.add_new_cross_reference()

        event_id = self.presenter.view.current_event.id
        self.assertFalse(event_id is None)

        # Assert that database has changed
        event_crossreferences_dao = self.injector.get(baseinjectorkeys.EVENT_CROSS_REFERENCES_DAO_KEY)
        cross_reference_ids = event_crossreferences_dao.get_cross_references(event_id)
        self.assertEqual(len(cross_reference_ids), len(reference_ids))
        for reference in reference_ids:
            self.assertIn(reference, cross_reference_ids)

        # Assert that view has been updated
        self.assertEqual(len(self.presenter.view.items), len(reference_ids))
        for reference in reference_ids:
            self.assertIn(self.event_dao.get_by_id(reference), self.presenter.view.items)

    def testAddNewCrossReferenceCancel(self):

        self.set_current_event(1940000001)
        self.presenter.view.new_cross_reference_event = None
        self.presenter.add_new_cross_reference()

        # Assert that database has not changed
        event_crossreferences_dao = self.injector.get(baseinjectorkeys.EVENT_CROSS_REFERENCES_DAO_KEY)
        cross_reference_ids = event_crossreferences_dao.get_cross_references(1940000001)
        self.assertEqual(len(cross_reference_ids), 2)
        self.assertIn(1950000001, cross_reference_ids)
        self.assertIn(1960013001, cross_reference_ids)

        # Assert that view has not changed
        self.assertEqual(len(self.presenter.view.items), 2)
        self.assertIn(self.event_dao.get_by_id(1950000001), self.presenter.view.items)
        self.assertIn(self.event_dao.get_by_id(1960013001), self.presenter.view.items)

    def testDeleteCrossReference(self):

        self.set_current_event(1940000001)
        self.presenter.view.selected_item = self.event_dao.get_by_id(1960013001)
        self.presenter.delete_cross_reference()

        # Assert that database has changed
        event_crossreferences_dao = self.injector.get(baseinjectorkeys.EVENT_CROSS_REFERENCES_DAO_KEY)
        cross_reference_ids = event_crossreferences_dao.get_cross_references(1940000001)
        self.assertEqual(len(cross_reference_ids), 1)
        self.assertEqual(cross_reference_ids[0], 1950000001)

        # Assert that view has been updated
        self.assertEqual(len(self.presenter.view.items), 1)
        self.assertEqual(self.presenter.view.items[0], self.event_dao.get_by_id(1950000001))

    def testDeleteCrossReferenceNothingSelected(self):

        self.set_current_event(1940000001)
        self.presenter._update_crossreferences() # Initialize items
        self.presenter.view.selected_item = None # If we wouldn't intialize, we'd get a mock

        self.presenter.delete_cross_reference()

        # Assert that database has not changed
        event_crossreferences_dao = self.injector.get(baseinjectorkeys.EVENT_CROSS_REFERENCES_DAO_KEY)
        cross_reference_ids = event_crossreferences_dao.get_cross_references(1940000001)

        self.assertEqual(len(cross_reference_ids), 2)
        self.assertIn(1950000001, cross_reference_ids)
        self.assertIn(1960013001, cross_reference_ids)

        # Assert that view is still correct
        self.assertEqual(len(self.presenter.view.items), 2)
        self.assertIn(self.event_dao.get_by_id(1950000001), self.presenter.view.items)
        self.assertIn(self.event_dao.get_by_id(1960013001), self.presenter.view.items)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()


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

# read the data
data = pd.read_csv('data.csv')

# Convert the Resolution time from seconds to hours
data['Resolution time'] = data['Resolution time'] / 3600

# Set the figure size and style
plt.figure(figsize=[12, 8])
sns.set(style="whitegrid")

# Boxplot visualization
sns.boxplot(data['Resolution time'], color='skyblue', width=0.4)

# Adding stripplot with adjusted jitter
sns.stripplot(data['Resolution time'], color='black', size=4, jitter=0.2)

# Set y-axis range

# Set the title and labels
plt.title('Distribution of Resolution Time')
plt.xlabel('Resolution Time (in hours)')

plt.savefig('aos/5_resolution.png', dpi=300, bbox_inches='tight')


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider, Button, TextBox
from tkinter import messagebox as mbox

# Colors:
bg_color = "#4e5254"
hc_color = "#afb1b3"
hit_color = "#c75450"
main_color_1 = "#8dd3c7"
main_color_2 = "#71b1d0"

# Service parameters:
show_prev_track = False
hit_check = False

# Initial ratios:
angle_grad = 45  # °
efficiency = 30  # %
gunpowder = 10  # g/number
target_height, target_x, hit_y = 20, 450, 0  # m
g = 9.8  # m/s^2
x0, y0 = 0, 0  # starting point
q = 3800  # J/g
M = 100  # kg
m = 5  # kg

# Cannon const parameters:
r_wheel = 0.3
friction_coef = 0.45

# Calculated ratios:
angle = 0
initial_speed = 0
time_interval = 0
x_offset = 0
y_offset = 0

x_prev_data, y_prev_data = [], []

plt.style.use('dark_background')

fig = plt.figure()
fig.patch.set_facecolor("#3c3f41")
fig.suptitle("Cannon Firing Simulator", fontsize=16, fontweight="bold")
fig.canvas.manager.set_window_title("Cannon Firing Simulator")
error_window_title = "Cannon Firing Setup"

# Proportions:
ws = [1 / 6, 1 / 6, 1 / 3, 1 / 3]
hs = [3 / 28, 3 / 28, 1 / 28, 1 / 28, 1 / 28, 3 / 28, 3 / 28, 3 / 28, 3 / 28, 1 / 4]

gs = GridSpec(ncols=4, nrows=10, width_ratios=ws, height_ratios=hs, figure=fig)

# Main plane:
ax = plt.subplot(gs[:9, 2:], facecolor=bg_color)
ax.set_aspect("equal")

plt.xlabel("x, m")
plt.ylabel("y, m")
plt.grid(True)

# Bars:
ax_impulses = plt.subplot(gs[9, :2], facecolor=bg_color)
impulse_bar = ax_impulses.bar(["Projectile", "Cannon (1)", "Cannon (2)"], [0, 0, 0], color=main_color_1)
plt.xlabel("Impulses, kg*m/s")
plt.grid(True)

ax_forces = plt.subplot(gs[9, 2], facecolor=bg_color)
force_bar = ax_forces.bar(["Friction C", "Reaction C", "Gravity P"], [0, 0, 0], color=main_color_1)
plt.xlabel("Forces, N")
plt.grid(True)

ax_velocities = plt.subplot(gs[9, 3], facecolor=bg_color)
velocity_bar = ax_velocities.bar(["Initial P", "X-coordinate P", "Initial C"], [0, 0, 0], color=main_color_1)
plt.xlabel("Velocities, m/s")
plt.grid(True)

axButton_launch = plt.subplot(gs[0, :2])
button_launch = Button(ax=axButton_launch, label="Fire", color=bg_color, hovercolor=hc_color)

# Values computing:

def compute_x(t, track="bullet"):
    global x_offset, angle, initial_speed, friction_coef, r_wheel, m, M, g

    if track == "bullet":
        return x_offset + np.cos(angle) * initial_speed * t
    elif track == "cannon":
        deceleration = compute_force(force="friction") / M
        if np.cos(angle) * initial_speed * (m / M) > deceleration * t:
            return x_offset - np.cos(angle) * initial_speed * (m / M) * t + deceleration * (t ** 2) / 2
        else:
            return x_offset - ((np.cos(angle) * initial_speed * (m / M)) ** 2) / (2 * deceleration)

def compute_y(t, track="bullet"):
    global y_offset, angle, initial_speed, g

    if track == "bullet":
        return y_offset + np.sin(angle) * initial_speed * t - g * (t ** 2) / 2
    elif track == "cannon":
        return y_offset

def compute_force(mass="m", force="friction"):
    global friction_coef, m, M, g, r_wheel

    if force == "friction":
        return friction_coef * M * g / r_wheel
    elif force == "gravity":
        if mass == "m":
            return m * g
        elif mass == "M":
            return M * g

def compute_impulse(t=0, track="bullet"):
    global initial_speed, friction_coef, M, m, g

    if track == "bullet":
        return initial_speed * m
    elif track == "cannon":
        return max(0.0, initial_speed * m - compute_force(force="friction") * t)

def clear_track():
    global x_prev_data, y_prev_data

    for line in ax.get_lines():
        line.remove()

    if show_prev_track:
        ax.plot(x_prev_data, y_prev_data, "--", color=hc_color, lw=3)

# Initial computing:

def update_config():
    global angle, initial_speed, time_interval, x_offset, y_offset, target_height, target_x, hit_y

    angle = angle_grad * np.pi / 180
    initial_speed = np.sqrt((2 * efficiency * q * gunpowder) / (100 * m))
    time_interval = (np.sin(angle) * initial_speed + np.sqrt(
        (np.sin(angle) * initial_speed) ** 2 + 2 * g * y0)) / g
    x_offset = x0
    y_offset = y0

    if compute_x(time_interval, track="bullet") >= target_x:
        t_hit = (target_x - x_offset) / (np.cos(angle) * initial_speed)
        hit_y = compute_y(t_hit, track="bullet")

    return np.arange(0, time_interval, 0.1)

# Target plotting:

def plot_target():
    global target_x, target_height

    ax.plot([target_x, target_x], [0, target_height], color=main_color_1, lw=3)

# Main measures bars plotting:

def plot_bars():
    global time_interval, initial_speed, angle

    y_impulses_data = [compute_impulse(track="bullet"), compute_impulse(track="cannon"),
                       compute_impulse(t=time_interval, track="cannon")]

    idx = 0
    for rect in impulse_bar:
        rect.set_height(y_impulses_data[idx])
        idx += 1

    ax_impulses.relim()
    ax_impulses.autoscale_view(scalex=True, scaley=True)

    y_forces_data = [compute_force(force="friction"), compute_force(mass="M", force="gravity"),
                     compute_force(mass="m", force="gravity")]

    idx = 0
    for rect in force_bar:
        rect.set_height(y_forces_data[idx])
        idx += 1

    ax_forces.relim()
    ax_forces.autoscale_view(scalex=True, scaley=True)

    y_velocities_data = [initial_speed, np.cos(angle) * initial_speed, np.cos(angle) * initial_speed * (m / M)]

    idx = 0
    for rect in velocity_bar:
        rect.set_height(y_velocities_data[idx])
        idx += 1

    ax_velocities.relim()
    ax_velocities.autoscale_view(scalex=True, scaley=True)

# Main (cannon & cannonball) tracks animation:

def run_animation():
    global x_prev_data, y_prev_data

    def update_track(t):
        global target_x, target_height, hit_y, hit_check

        x_bullet = compute_x(t, track="bullet")
        y_bullet = compute_y(t, track="bullet")

        x_cannon = compute_x(t, track="cannon")
        y_cannon = compute_y(t, track="cannon")

        if x_bullet not in x_data["bullet"]:
            x_data["bullet"].append(x_bullet)
            y_data["bullet"].append(y_bullet)

            if target_x <= x_bullet and hit_y <= target_height and hit_check:
                ax.plot([target_x], [hit_y], "o", mfc=hit_color, mec=hit_color, markersize=8)
                hit_check = False

            bullet_track.set_data(x_data["bullet"], y_data["bullet"])

        if x_cannon not in x_data["cannon"]:
            x_data["cannon"].append(x_cannon)
            y_data["cannon"].append(y_cannon)

            cannon_track.set_data(x_data["cannon"], y_data["cannon"])

        ax.relim()
        ax.autoscale_view(scalex=True, scaley=True)

    bullet_track, = ax.plot([], [], color=main_color_2, lw=3)
    cannon_track, = ax.plot([], [], color=main_color_1, lw=3)

    x_data, y_data = {"cannon": [], "bullet": []}, {"cannon": [], "bullet": []}

    x_prev_data = x_data["bullet"]
    y_prev_data = y_data["bullet"]

    FuncAnimation(fig, func=update_track, frames=update_config(), interval=20, blit=False)

def launch(event):
    global hit_check

    hit_check = True
    clear_track()
    plot_target()
    plot_bars()
    run_animation()

button_launch.on_clicked(launch)

# Secondary widgets:

axButton_prev = plt.subplot(gs[1, :2])
button_prev = Button(ax=axButton_prev, label="Save previous track", color=bg_color, hovercolor=hc_color)

def update_prev(event):
    global show_prev_track

    if show_prev_track:
        show_prev_track = False
        button_prev.label.set_text("Save previous track")
    else:
        show_prev_track = True
        button_prev.label.set_text("Omit previous track")

button_prev.on_clicked(update_prev)

axSlider_angle = plt.subplot(gs[2, :2])
slider_angle = Slider(ax=axSlider_angle, label="Angle, °", valmin=1, valmax=89, valinit=45, valstep=1, initcolor=None,
                      track_color=hc_color)

def update_angle(val):
    global angle_grad
    angle_grad = slider_angle.val

slider_angle.on_changed(update_angle)

axSlider_gunpowder = plt.subplot(gs[3, :2])
slider_gunpowder = Slider(ax=axSlider_gunpowder, label="Gunpowder, g", valmin=10, valmax=100, valstep=10,
                          initcolor=None,
                          track_color=hc_color)

def update_gunpowder(val):
    global gunpowder
    gunpowder = slider_gunpowder.val

slider_gunpowder.on_changed(update_gunpowder)

axSlider_efficiency = plt.subplot(gs[4, :2])
slider_efficiency = Slider(ax=axSlider_efficiency, label="Efficiency, %", valmin=1, valmax=100, valstep=1, valinit=30,
                           initcolor=None, track_color=hc_color)

def update_efficiency(val):
    global efficiency
    efficiency = slider_efficiency.val

slider_efficiency.on_changed(update_efficiency)

# Error handling:

def check_format(label, positive_only=False):
    try:
        value = int(label)
        if (not positive_only and value >= 0) or value > 0:
            return True
        else:
            mbox.showerror(error_window_title, "Incorrect measure format")
            return False

    except ValueError:
        mbox.showerror(error_window_title, "Incorrect measure format")
        return False

axTextBox_x_start = plt.subplot(gs[5, 0])
axTextBox_x_start.set_title("Starting point X")

textbox_x_start = TextBox(ax=axTextBox_x_start, label="X, m", initial="0", textalignment="center", color=bg_color,
                          hovercolor=hc_color)

def update_x0(label):
    global x0

    if check_format(label):
        x0 = int(label)

textbox_x_start.on_submit(update_x0)

axTextBox_y_start = plt.subplot(gs[5, 1])
axTextBox_y_start.set_title("Starting point Y")

textbox_y_start = TextBox(ax=axTextBox_y_start, label="Y, m", initial="0", textalignment="center", color=bg_color,
                          hovercolor=hc_color)

def update_y0(label):
    global y0

    if check_format(label):
        y0 = int(label)

textbox_y_start.on_submit(update_y0)

axTextBox_target_x = plt.subplot(gs[6, 0])
axTextBox_target_x.set_title("Target X")

textbox_target_x = TextBox(ax=axTextBox_target_x, label="X, m", initial="450", textalignment="center", color=bg_color,
                           hovercolor=hc_color)

def update_target_x(label):
    global target_x

    if check_format(label):
        target_x = int(label)

textbox_target_x.on_submit(update_target_x)

axTextBox_target_height = plt.subplot(gs[6, 1])
axTextBox_target_height.set_title("Target height")

textbox_target_height = TextBox(ax=axTextBox_target_height, label="H, m", initial="20", textalignment="center",
                                color=bg_color, hovercolor=hc_color)

def update_target_height(label):
    global target_height

    if check_format(label):
        target_height = int(label)

textbox_target_height.on_submit(update_target_height)

axTextBox_bullet_m = plt.subplot(gs[7, :2])
axTextBox_bullet_m.set_title("Projectile mass:")

textbox_bullet_m = TextBox(ax=axTextBox_bullet_m, label="m, kg", initial="5", textalignment="center", color=bg_color,
                           hovercolor=hc_color)

def update_bullet_m(label):
    global m

    if check_format(label, positive_only=True):
        m = int(label)

textbox_bullet_m.on_submit(update_bullet_m)

axTextBox_cannon_m = plt.subplot(gs[8, :2])
axTextBox_cannon_m.set_title("Cannon mass:")

textbox_cannon_m = TextBox(ax=axTextBox_cannon_m, label="M, kg", initial="100", textalignment="center", color=bg_color,
                           hovercolor=hc_color)

def update_cannon_m(label):
    global M

    if check_format(label, positive_only=True):
        M = int(label)

textbox_cannon_m.on_submit(update_cannon_m)

plt.subplots_adjust(hspace=1)
plt.show()


#!/usr/bin/env python3

import psutil
import re
import time
import os.path
import datetime
import subprocess

def find_proc_by_regex(regex, username):
    "Return a list of processes matching 'cmdline' and 'username'."
    my_regex = re.compile(regex)
    for p in psutil.process_iter(['cmdline', 'username', 'pid']):
            if (my_regex.search(" ".join(p.info['cmdline'])) is not None) and \
               (p.info['username'] == username):
                return p

def time_used(process):
    time_now = time.time()
    process_start = process.create_time()
    used =  time_now - process_start

    print("time now: {}".format(time_now))
    print("process started: {}".format(process_start))
    print("seconds used by this instance: {}".format(used))
    return used

def track_process(process, timebudget, grace, time_granularity):
    print("Entering track_process(), time budget is {}".format(timebudget))
    used = 0
    while process.is_running():
        print("Tracking process {}".format(process))
        used = time_used(process)
        if timebudget - 60 <= used:
            subprocess.run([ "su", "-", "ian", "-c", "paplay /usr/share/sounds/ubuntu/ringtones/Harmonics.ogg"])
        if timebudget <= used:
            print("Terminating target process")
            process.terminate()
            time.sleep(grace)
            if process.is_running():
                print("Killing target process")
                process.kill()
            print("Process terminated/killed")
        time.sleep(time_granularity)
    return used

def main(name, regex, username, timelimit, grace, time_granularity):
    filepath = os.path.join('/tmp', name + "_seconds_used")
    while True:
        print("Checking for matching process")
        process = find_proc_by_regex(regex, username)
        if process:
            used = 0
            if os.path.isfile(filepath):
                # check that timestamp is from today
                if datetime.datetime.fromtimestamp(os.path.getmtime(filepath)).day == \
                    datetime.date.today().day:
                        with open(filepath) as f:
                            try:
                                used = float(f.read())
                            except:
                                pass
            used += track_process(process, timelimit - used, grace, time_granularity)
            print("Process exited after using {} seconds".format(used))
            with open(os.path.join('/tmp', name + "_seconds_used"), 'w+') as f:
                f.write(str(used))
        else:
            print("No matching process found")

        time.sleep(time_granularity)

if __name__ == "__main__":

    name = 'minecraft'
    regex = 'java.*minecraft'
    username = 'ian'
    #timelimit = 60 * 60 * 3
    timelimit = 60 * 3
    grace = 5
    time_granularity = 10

    main(name, regex, username, timelimit, grace, time_granularity)


from django.shortcuts import redirect, render
from item.models import Item, Category
from django.contrib.auth.models import auth

from . forms import SignUpForm
from django.contrib import messages

def frontpage(request):
    items=Item.objects.filter(is_sold=False)[0:4]
    categories=Category.objects.all()
    context={
        'items': items,
        'categories': categories
    }
    return render(request, 'core/frontpage.html', context)

def contact(request):
    return render(request, 'core/contact.html')

def signup(request):
    if request.method=='POST':
        form=SignUpForm(request.POST)

        if form.is_valid():
            form.save()
            messages.success(request, 'Вы успешно зарегистрировались!')
            return redirect('/login/')
    else:
        form=SignUpForm()
    context={
        'form': form
    }
    return render(request, 'core/signup.html', context)

def logout(request):
    auth.logout(request)
    messages.success(request, 'Вы успешно вышли из системы')
    return redirect('core:login')


import logging
import socket

from support.process import invoke

module_logger = logging.getLogger(__name__)

# I need to get a definitive list for each domain
networks = {"fltops": ["137.228.202",
                       "137.228.203",
                       "137.228.207",
                       "137.228.236",
                       "137.228.246",
                       "137.228.247"],
           "jpl": ["128.149.22",
                   "128.149.252",
                   "137.79.89"]}

services = {"rdp": 3389, "vnc": 5900, "vncJava": 5800, "ssh": 22, "MSrcp": 135,
            "ftp": 21, "sftp": 115}

def LAN_hosts_status(port=22, private=False):
  """
  Get information about and status of local network hosts with port open

  Returns a dict with::
    - "up":    hosts up
    - "down":  hosts down
    - "IP":    IP addresses
    - "MAC":   MAC addresses
    - "ROACH": list of ROACHes

  This is too crude; very senstive to format changes
  """
  print("If asked for a password, remember this host is",socket.gethostname())
  domain = get_local_network(private=private)
  response = invoke("sudo nmap -sP -PS"+str(port)+" "+domain+".*").stdout.readlines()
  up = []
  down = []
  IP = {}
  MAC = {}
  ROACHlist = []
  for line in range(4,len(response)-1):
    parts = response[line].strip().split()
    if parts[2] == "report":
      host = parts[4]
      if len(parts) == 6:
        IP[host] = parts[5][1:-1]
      else:
        IP[host] = host
      line += 1
      status = response[line].strip().split()[2]
      if status == 'up':
        up.append(host)
        line += 1
        MAC[host] = response[line].strip().split()[2]
        if MAC[host].split(":")[0] == '02':
          ROACHlist.append(host)
      else:
        down.append(host)
    else:
      pass
  up.sort()
  down.sort()
  ROACHlist.sort()
  return {"up": up,
          "down": down,
          "IP": IP,
          "MAC:": MAC,
          "ROACH": ROACHlist}

def report_LAN_hosts_status():
  up, down, IP, MAC, ROACHlist = LAN_hosts_status()
  return {"up": up,
          "down": down,
          "IP": IP,
          "MAC": MAC,
          "roaches detected": ROACHlist}

def get_local_network(private=False, subnet=100):
  """
  Returns the IP address of the local network

  Defaults to the local area network.  If 'private", defaults to '100' for
  the ROACH network.
  """
  if private:
    IP = socket.gethostbyname("192.168."+str(subnet)+".1")
  else:
    IP = socket.gethostbyname(socket.gethostname())
  return '.'.join(IP.split('.')[:-1])

def get_domain(netIP):
  """
  Returns the networks key whose list contains netIP
  """
  domain = ""
  for key in list(networks.keys()):
    for IP in networks[key]:
      if IP == netIP:
        domain = key
        break
  if domain == "":
    domain = "other"
  return domain

################################ IP support ###################################

def decode_IP(IP_address):
  """
  Returns a long int for an IPv4 or IPv6 address.

  @param IP_address :  like "192.168.1.101"
  @type  IP_address : str of decimal separated ints

  @return: int
  """
  parts = IP_address.split('.')
  if len(parts) == 4 or len(parts) == 6:
    ipvalue = 0
    for i in range(len(parts)):
      shift = 8*(len(parts)-1-i)
      ipvalue += int(parts[i]) << shift
    return ipvalue
  else:
    raise RuntimeError("Invalid IP address: %s" % IP_address)

def decode_MAC(MAC_address):
  """
  Returns a long int for a MAC address

  @param MAC_address : like "00:12:34:56:78:9a"
  @type  MAC_address : str of colon separated hexadecimal ints

  @return: long int
  """
  parts = MAC_address.split(':')
  if len(parts) == 6:
    value = 0
    for i in range(6):
      shift = 8*(5-i)
      value += int(parts[i],16) << shift
    return value
  else:
    raise RuntimeError("Invalid MAC address: %s" % MAC_address)

def encode_IP(intIP):
  """
  """
  part0 = intIP/2**24
  part1 = (intIP - part0*2**24)/2**16
  part2 = (intIP - part0*2**24 - part1*2**16)/2**8
  part3 =  intIP - part0*2**24 - part1*2**16 - part2*2**8
  return "%d.%d.%d.%d" % (part0, part1, part2, part3)

if __name__ == "__main__":
  up, down, IP, MAC, ROACHlist = LAN_hosts_status()
  print("Up:", up)
  print("Down:", down)
  print("ROACHes:", ROACHlist)
  print("IP:", IP)
  print("MAC addresses:", MAC)


from manim import Scene, Graph, Create

class Anim(Scene):
    def construct(self):
        vertices = [1, 2]
        edges = [(1, 2)]

        layout = {1: [-1, 0, 1], 2: [1, 0, 2]}

        graph = Graph(
            vertices,
            edges,
            layout=layout,
            layout_scale=3,
        )
        self.play(Create(graph))
        self.wait(1)


"""
Unit system for physical quantities; include definition of constants.
"""

from typing import Dict

from sympy import S, Mul, Pow, Add, Function, Derivative
from sympy.physics.units.dimensions import _QuantityMapper

from sympy.utilities.exceptions import SymPyDeprecationWarning

from .dimensions import Dimension

class UnitSystem(_QuantityMapper):
    """
    UnitSystem represents a coherent set of units.

    A unit system is basically a dimension system with notions of scales. Many
    of the methods are defined in the same way.

    It is much better if all base units have a symbol.
    """

    _unit_systems = {}  # type: Dict[str, UnitSystem]

    def __init__(self, base_units, units=(), name="", descr="", dimension_system=None):

        UnitSystem._unit_systems[name] = self

        self.name = name
        self.descr = descr

        self._base_units = base_units
        self._dimension_system = dimension_system
        self._units = tuple(set(base_units) | set(units))
        self._base_units = tuple(base_units)

        super().__init__()

    def __str__(self):
        """
        Return the name of the system.

        If it does not exist, then it makes a list of symbols (or names) of
        the base dimensions.
        """

        if self.name != "":
            return self.name
        else:
            return "UnitSystem((%s))" % ", ".join(
                str(d) for d in self._base_units)

    def __repr__(self):
        return '<UnitSystem: %s>' % repr(self._base_units)

    def extend(self, base, units=(), name="", description="", dimension_system=None):
        """Extend the current system into a new one.

        Take the base and normal units of the current system to merge
        them to the base and normal units given in argument.
        If not provided, name and description are overridden by empty strings.
        """

        base = self._base_units + tuple(base)
        units = self._units + tuple(units)

        return UnitSystem(base, units, name, description, dimension_system)

    def print_unit_base(self, unit):
        """
        Useless method.

        DO NOT USE, use instead ``convert_to``.

        Give the string expression of a unit in term of the basis.

        Units are displayed by decreasing power.
        """
        SymPyDeprecationWarning(
            deprecated_since_version="1.2",
            issue=13336,
            feature="print_unit_base",
            useinstead="convert_to",
        ).warn()
        from sympy.physics.units import convert_to
        return convert_to(unit, self._base_units)

    def get_dimension_system(self):
        return self._dimension_system

    def get_quantity_dimension(self, unit):
        qdm = self.get_dimension_system()._quantity_dimension_map
        if unit in qdm:
            return qdm[unit]
        return super().get_quantity_dimension(unit)

    def get_quantity_scale_factor(self, unit):
        qsfm = self.get_dimension_system()._quantity_scale_factors
        if unit in qsfm:
            return qsfm[unit]
        return super().get_quantity_scale_factor(unit)

    @staticmethod
    def get_unit_system(unit_system):
        if isinstance(unit_system, UnitSystem):
            return unit_system

        if unit_system not in UnitSystem._unit_systems:
            raise ValueError(
                "Unit system is not supported. Currently"
                "supported unit systems are {}".format(
                    ", ".join(sorted(UnitSystem._unit_systems))
                )
            )

        return UnitSystem._unit_systems[unit_system]

    @staticmethod
    def get_default_unit_system():
        return UnitSystem._unit_systems["SI"]

    @property
    def dim(self):
        """
        Give the dimension of the system.

        That is return the number of units forming the basis.
        """
        return len(self._base_units)

    @property
    def is_consistent(self):
        """
        Check if the underlying dimension system is consistent.
        """
        # test is performed in DimensionSystem
        return self.get_dimension_system().is_consistent

    def get_dimensional_expr(self, expr):
        from sympy import Mul, Add, Pow, Derivative
        from sympy import Function
        from sympy.physics.units import Quantity
        if isinstance(expr, Mul):
            return Mul(*[self.get_dimensional_expr(i) for i in expr.args])
        elif isinstance(expr, Pow):
            return self.get_dimensional_expr(expr.base) ** expr.exp
        elif isinstance(expr, Add):
            return self.get_dimensional_expr(expr.args[0])
        elif isinstance(expr, Derivative):
            dim = self.get_dimensional_expr(expr.expr)
            for independent, count in expr.variable_count:
                dim /= self.get_dimensional_expr(independent)**count
            return dim
        elif isinstance(expr, Function):
            args = [self.get_dimensional_expr(arg) for arg in expr.args]
            if all(i == 1 for i in args):
                return S.One
            return expr.func(*args)
        elif isinstance(expr, Quantity):
            return self.get_quantity_dimension(expr).name
        return S.One

    def _collect_factor_and_dimension(self, expr):
        """
        Return tuple with scale factor expression and dimension expression.
        """
        from sympy.physics.units import Quantity
        if isinstance(expr, Quantity):
            return expr.scale_factor, expr.dimension
        elif isinstance(expr, Mul):
            factor = 1
            dimension = Dimension(1)
            for arg in expr.args:
                arg_factor, arg_dim = self._collect_factor_and_dimension(arg)
                factor *= arg_factor
                dimension *= arg_dim
            return factor, dimension
        elif isinstance(expr, Pow):
            factor, dim = self._collect_factor_and_dimension(expr.base)
            exp_factor, exp_dim = self._collect_factor_and_dimension(expr.exp)
            if exp_dim.is_dimensionless:
                exp_dim = 1
            return factor ** exp_factor, dim ** (exp_factor * exp_dim)
        elif isinstance(expr, Add):
            factor, dim = self._collect_factor_and_dimension(expr.args[0])
            for addend in expr.args[1:]:
                addend_factor, addend_dim = \
                    self._collect_factor_and_dimension(addend)
                if dim != addend_dim:
                    raise ValueError(
                        'Dimension of "{}" is {}, '
                        'but it should be {}'.format(
                            addend, addend_dim, dim))
                factor += addend_factor
            return factor, dim
        elif isinstance(expr, Derivative):
            factor, dim = self._collect_factor_and_dimension(expr.args[0])
            for independent, count in expr.variable_count:
                ifactor, idim = self._collect_factor_and_dimension(independent)
                factor /= ifactor**count
                dim /= idim**count
            return factor, dim
        elif isinstance(expr, Function):
            fds = [self._collect_factor_and_dimension(
                arg) for arg in expr.args]
            return (expr.func(*(f[0] for f in fds)),
                    expr.func(*(d[1] for d in fds)))
        elif isinstance(expr, Dimension):
            return 1, expr
        else:
            return expr, Dimension(1)


'''
Condition
    - TL : 1s (2000만)
      ML : 256 (64*100만)
    - 1<=N,K<=10만
    - -1억<=샘터 위치<=1억
    - 불행도 : 집과, 집에서 가장 가까운 샘터까지의 거리

Question
    - 샘터의 위치가 고정되어 있을 때 ,모든 집에 대한 불행도 합의 최솟값을 출력하여라

Access
    try_1)
        - 예상 분류 : dp
          문제 분류 : bfs

        - 만약 dp라면?
            >>> dp[k]= k번째 집이 있을 때 불행도 합, 모든 position에 넣어보고 최솟값을 가져온다
                : position 범위가 너무 크다..!!

        - bfs
            >>> 각 샘터를 bfs 시작지점으로 두고, 하나씩 채워나가면 된다!
                1. bfs 시작점을 샘터 시작점으로 잡는다
                2. 하나의 샘터 양쪽으로 집을 짓고, 그 다음 샘터로 넘어가서 양쪽으로 집을 짓는다
                3. 샘터의 개수만큼 2번을 반복한다
                4. 집의 개수만큼 2~3번을 반복한다

            문제점
                1. 각 집에서 가장 가까운 샘터를 어떻게 알아낼 것인가?
                    - 각 집마다 모든 샘터를 비교 >> n*k = 최대 10만*10만 = TLE 발생
                    - 각 집에서 bfs를 시전해서 가까운 샘터를 찾아야한다 !!! >> 집위치가 샘터와 멀어지면 무조건 TLE

                >>> bfs에서 dq에 집의 위치와 거리를 동시에 포함시키면 해결된다

'''
from collections import deque
from sys import stdin
input = stdin.readline

#define function
def doVisit(position,dist):
    global cnt
    visited.add(position)       #방문처리
    possible.append((position,dist+1))   #다음건설위치 추가
    cnt +=1                     #건설된 집의 수 증가

def solution():
    global cnt, visited, possible
    cnt = 0                   # 건설된 집의 수
    unhappiness = 0           # 불행도
    visited = set(springSite) # 건설된 집이나 샘의 위치를 방문처리

    # possible은 집 위치 후보로, (위치,거리)를 포함한다
    possible = deque()
    for spring in springSite:
        possible.append((spring,1))

    # 건설된 집의 수가 k개가 될때까지 반복
    while True:
        posit,dist = possible.popleft()
        positRight = posit +1
        positLeft = posit -1

        # posit 오른편에 집을 지음 : 방문처리 + 다음 건설위치에 추가 + 건설된 집의 수 증가 + 불행도 갱신
        if positRight not in visited:
            doVisit(positRight,dist)
            unhappiness += dist # 불행도 갱신

        # 건설된 집의 개수가 k개라면 종료
        if cnt == k:
            return unhappiness

        # posit 왼쪽편에 집을 지음
        if positLeft not in visited:
            doVisit(positLeft,dist)
            unhappiness += dist # 불행도 갱신

        # 건설된 집의 개수가 k개라면 종료
        if cnt == k:
            return unhappiness

#main
n,k = map(int,input().split())
springSite = list(map(int,input().split())) # 샘터 위치
print(solution())


# Keiland Pullen
# DSC 430: Python Programming
# Due Date: 11/26/2019
# Link to video: https://youtu.be/8tKA-NuqJJw
# Link to Second Chance video: https://youtu.be/FAcugzuUA6o
# Honor Statement: I have not given or received any unauthorized assistance on this assignment
# Assignment 1001: Second Chance

from bs4 import BeautifulSoup
from urllib.request import urlopen
import requests
from collections import Counter
from urllib.parse import urlparse
from urllib.parse import  urljoin

def analyze(url):
    '''Function to return Visited URL and word frequency on that page along with list of http links at that url.'''

    print('\n \nVisiting', url)

    getTheURL = requests.get(url)                       #get the URL
    data = getTheURL.text                               #capture the text from that URL
    soup = BeautifulSoup(data, 'html.parser')           #user BeautifulSoup module to parse the text
    urls = []                                           #initiate list of URLs
    wordList = []                                       #initiate list of words

    for link in soup.find_all('a'):                     #BeautifulSoup module to find all anchor links on page

        urls.append(urljoin(url,link.get('href')))      #for each link, create absolute URL

        for allText in soup.find_all('body'):           #get all text within the BODY tag
            content = allText.text
            words = content.lower().split()             #filter each word
            for eachWord in words:
                wordList.append(eachWord)               #add each filtered word into list

    print('\n{:45} {:10} {:5}'.format('URL', 'Word', 'Count'))      # print header

    washWords(wordList, url)                            #send words to be "cleaned".

    print('\n{:45} {:10}'.format('URL', 'Link'))

    for link in urls:
        print('{:45} {:10}'.format(url, link))          #print each parent url and the link on its page

    return urls

def washWords(myWords, url):
    '''Function to remove non-alpha numeric characters from the word list. '''

    url = url
    cleanWordsList = []

    for word in myWords:
        chars = '!@#$%^&*()_-+={[}]|\;:"<>?/., '

        for x in range (0, len(chars)):
            word = word.replace(chars[x], '')

        if len(word) > 0:
            cleanWordsList.append(word)

    wordDict(cleanWordsList, url)

def wordDict(cleanWordsList, url):
    '''Function to display the 5 most frequent words on a page  '''

    wordCount = {}

    for word in cleanWordsList:
        if word in wordCount:
            wordCount[word] += 1
        else:
            wordCount[word] = 1

    counts = Counter(wordCount)

    freq = counts.most_common(5)                       #the frequency of the top 5 words on a page
    #print(freq)

    for key, value in freq:
        #print (url, key, value)                        # print url, key, value for testing
        print('\n{:45} {:10} {:5}'.format(url, key, value))

visited = set()                                         #set to contain visited URLs

def crawl(url):
    '''Recursive web crawler'''

    global visited
    visited.add(url)
    #print("Vistited links and urls: ",visited)          # print for testing

    links = analyze(url)

    for link in links:
        #print(link)                                    # print links -- for testing
        if link not in visited:
            #try:                                       # try except in testing
            parsedURL = urlparse(link)
            #print(parsedURL.netloc)                    # print URL -- for testing

            cdmPath = parsedURL.netloc

            if cdmPath == "cdm.depaul.edu":
                crawl(link)
            #except:                                    # try except in testing
            #    pass


"""This file creates a trainer that finds the best (fittest) AI"""
from ai_player import AiPlayer
from game import Game, SampleStrategies
import numpy as np
from numpy import random
import random
from typing import List, Tuple, Any
import pickle
from copy import copy, deepcopy

class AiLogEntry:
    """data to track evolution of AIs across generations during training"""

    def __init__(self, points_average, points_variance, events):
        self.points_average = points_average
        self.points_variance = points_variance
        self.events = events

    def __repr__(self):
        return "(" + str(self.points_average) + ", " + str(self.points_variance) + ", " + str(self.events) + ")"

class Trainer:
    """trains AIs (with a genetic algorithm) in order to get the best AI out of all;
    parameters therefore given manually"""

    strategy_parameter_min = -10
    strategy_parameter_max = 10

    def __init__(self, group_size=5, play_against_own_copies=False, population_size=100, survivor_rate=0.74,
                 child_rate=1, mutation_rate=0.05, mutation_copy_rate=0.0, lowest_variance_rate=0.98,
                 num_generations=200):
        assert (population_size % group_size == 0)
        self.group_size = group_size
        self.play_against_own_copies = play_against_own_copies
        self.population_size = population_size
        self.mutation_rate = mutation_rate
        self.mutation_copy_rate = mutation_copy_rate
        self.lowest_variance_rate = lowest_variance_rate
        self.num_generations = num_generations
        self.survivor_rate = survivor_rate
        self.num_survivors = int(survivor_rate * self.population_size)
        self.child_rate = child_rate
        self.num_children = int((self.population_size - self.num_survivors) * child_rate) # child_rate is relative amount of died AIs that is 'reborn' by recombination
        self.num_parents = self.num_children * 2
        self.fitness_game_number = 5   # empiric value
        self.points_per_ai = None
        self.ai_histories = dict()
        self.generation = 0
        self.population = []

    def _group(self) -> List[Tuple[Any, List[AiPlayer]]]:
        """puts the individuals (AIs) of a population into random groups of the same size"""
        if self.play_against_own_copies:
            groups = [(ai, [deepcopy(ai) for _ in range(self.group_size)]) for ai in self.population]
        else:
            random.shuffle(self.population)
            groups = []
            for ai_index in range(0, len(self.population), self.group_size):
                groups.append(("different AIs", self.population[ai_index: ai_index + self.group_size]))
        return groups

    def _select_extreme_ais(self, point_sum_per_ai, num_extreme_ais, selects_best_ais, variance_threshold) -> List[
        AiPlayer]:
        """selects strongest or weakest AIs of a population"""
        extreme_ais = []
        for _ in range(num_extreme_ais):
            max_points = float("-inf")
            min_points = float("inf")
            extreme_ai = None
            for ai, points in point_sum_per_ai.items():
                if self.ai_histories[ai][self.generation].points_variance > variance_threshold:
                    continue    # ignore AIs with high points variance in ranking -> middle field (extreme good/bad daily form)
                if selects_best_ais and max_points < points:
                    max_points = points
                    extreme_ai = ai
                elif not selects_best_ais and min_points > points:
                    min_points = points
                    extreme_ai = ai
            extreme_ais.append(extreme_ai)
            del point_sum_per_ai[extreme_ai]
        assert(None not in extreme_ais)
        return extreme_ais

    def _compute_avg_points_per_ai(self, point_list_per_ai) -> None:
        """creates a dictionary with average points of every AI"""
        self.points_per_ai = dict()
        for ai, point_list in point_list_per_ai.items():
            self.points_per_ai[ai] = sum(point_list) / len(point_list)
            points_variance = Trainer._compute_variance(point_list_per_ai[ai], self.points_per_ai[ai])
            if self.generation in self.ai_histories[ai]:
                self.ai_histories[ai][self.generation].points_average = int(self.points_per_ai[ai])
                self.ai_histories[ai][self.generation].points_variance = int(points_variance)
            else:
                self.ai_histories[ai][self.generation] = \
                    AiLogEntry(int(self.points_per_ai[ai]), int(points_variance), [])

    @staticmethod
    def _compute_variance(numbers, average) -> float:
        """calculates the corrected sample variance of any list of numbers with a given average"""
        variance = 0  # korrigierte Stichprobenvarianz
        for number in numbers:
            variance += (number - average) ** 2
        return variance / len(numbers)

    def _play_in_groups(self, point_list_per_ai, game_count) -> None:
        """part of _rank(); lets AIs play in groups"""
        groups = self._group()
        for original_ai, group in groups:
            game = Game(group)
            # same lists of dice eyes for each generation to make AIs more comparable
            random.seed(self.generation * self.fitness_game_number + game_count)
            game.play()
            for ai in group:
                points = ai.get_points()
                if original_ai == "different AIs":
                    point_list_per_ai[ai].append(points)
                else:
                    point_list_per_ai[original_ai].append(points)

    def _play_against_own_copies(self, final_ai):
        """lets final AI play against its copies"""
        group = [deepcopy(final_ai) for _ in range(self.group_size)]
        point_list = []
        for _ in range(10):
            game = Game(group)
            game.play()
            for ai in group:
                point_list.append(ai.get_points())
        average = sum(point_list) / len(point_list)
        return average, self._compute_variance(point_list, average)

    def _split_ais_by_fitness(self, point_sum_per_ai):
        """part of _rank(); splits AIs by fitness after playing in groups"""
        variance_threshold = self._calculate_variance_threshold()
        point_sum_per_ai_temp = copy(point_sum_per_ai)
        strongest_ais = self._select_extreme_ais(point_sum_per_ai_temp, self.num_parents,
                                                 True, variance_threshold)  # side-effect intentional
        weakest_ais = self._select_extreme_ais(point_sum_per_ai_temp, self.population_size - self.num_survivors,
                                               False, variance_threshold)
        weakest_ais = list(reversed(weakest_ais))
        middle_field = point_sum_per_ai_temp.keys()  # keys() only selects the AIs, not the points
        return strongest_ais, middle_field, weakest_ais

    def _calculate_variance_threshold(self):
        """calculates the threshold of variance; necessary to put AIs with extreme points variance into the middle field
        (influence from daily form too high)"""
        variances = []
        for ai in self.population:
            variances.append(self.ai_histories[ai][self.generation].points_variance)
        return list(sorted(variances))[int(self.lowest_variance_rate * self.population_size) - 1]

    def _create_ranking(self, point_list_per_ai):
        """part of _rank(); creates a ranking based on the fitness of an AI"""
        point_sum_per_ai = {ai: sum(point_list) for ai, point_list in point_list_per_ai.items()}
        strongest_ais, middle_field, weakest_ais = self._split_ais_by_fitness(point_sum_per_ai)
        ranking = copy(strongest_ais)
        ranking.extend(middle_field)
        ranking.extend(weakest_ais)
        assert (len(ranking) == self.population_size)
        return ranking

    def _rank(self) -> None:
        """lets the groups play and ranks them inside these groups by performance"""
        point_list_per_ai = {ai: [] for ai in self.population}
        for game_count in range(self.fitness_game_number):
            self._play_in_groups(point_list_per_ai, game_count)
        if not self.play_against_own_copies:
            for point_list in point_list_per_ai.values():
                assert (len(point_list) == self.fitness_game_number)
        self._compute_avg_points_per_ai(point_list_per_ai)
        self.population = self._create_ranking(point_list_per_ai)

    def _select(self) -> None:
        """selects the best AIs, these survive -> rate of survivors given by parameter"""
        dead_ais = self.population[self.num_survivors:]
        for ai in dead_ais:
            self._add_event_to_ai_history(ai, "SELEction")
        self.population = self.population[: self.num_survivors]

    def _mix_strategies(self, parent1, parent2) -> AiPlayer:
        """mixes the strategies of the parents by creating averages of the different factors / bias to be their child's
         strategy"""
        child_quadratic_factor = np.array([x / 2 for x in (parent1.quadratic_factor + parent2.quadratic_factor)])
        child_linear_factor = np.array([x / 2 for x in (parent1.linear_factor + parent2.linear_factor)])
        child_bias = np.array([x / 2 for x in (parent1.bias + parent2.bias)])
        assert (len(child_linear_factor) == len(parent1.linear_factor) and len(child_linear_factor) == len(
            parent2.linear_factor))
        self._add_event_to_ai_history(parent1, "PAREnt")
        self._add_event_to_ai_history(parent2, "PAREnt")
        ai_number = len(self.ai_histories)
        return AiPlayer(str(ai_number), child_quadratic_factor, child_linear_factor, child_bias)

    def _add_event_to_ai_history(self, ai, event) -> None:
        if self.generation in self.ai_histories[ai]:
            self.ai_histories[ai][self.generation].events.append(event)
        else:
            self.ai_histories[ai][self.generation] = AiLogEntry(None, None, [event])

    def _recombine(self) -> None:
        """extends the population by children that are created by recombination of their parents strategies"""
        children = []
        for child_index in range(self.num_children):  # build pairs
            child = self._mix_strategies(self.population[child_index * 2], self.population[child_index * 2 + 1])
            children.append(child)
            self.ai_histories[child] = dict()
            self._add_event_to_ai_history(child, "RECOmbination")
        self.population.extend(children)
        assert (len(self.population) <= self.population_size)

    def _mutate_strategy(self, ai) -> int:
        """mutates randomly small parts of the strategy of an AI"""
        mutation_counter = 0
        for value_index in range(AiPlayer.strategy_length):
            if random.random() < self.mutation_rate:
                ai.quadratic_factor[value_index] = Trainer._adjust_strategy_range(random.random())
                mutation_counter += 1
            if random.random() < self.mutation_rate:
                ai.linear_factor[value_index] = Trainer._adjust_strategy_range(random.random())
                mutation_counter += 1
            if random.random() < self.mutation_rate:
                ai.bias[value_index] = Trainer._adjust_strategy_range(random.random())
                mutation_counter += 1
        return mutation_counter  # call by reference on ai (pointer)

    def _mutate(self) -> None:
        """creates a list of mutated AIs"""
        copied_ais = []
        max_copies = (self.population_size - len(self.population)) * self.mutation_copy_rate
        for ai in self.population:
            ai_copy = deepcopy(ai)
            ai_number = len(self.ai_histories)
            ai_copy.name = str(ai_number)
            ai_history_copy = deepcopy(self.ai_histories[ai])
            mutation_counter = self._mutate_strategy(ai)
            if mutation_counter > 0:
                self._add_event_to_ai_history(ai, f"MUTAtion, {mutation_counter}")
                if len(copied_ais) < max_copies:
                    copied_ais.append(ai_copy)
                    self.ai_histories[ai_copy] = ai_history_copy
                    self._add_event_to_ai_history(ai_copy, f"COPY from {ai}")
        self.population.extend(copied_ais)
        assert (len(self.population) <= self.population_size)

    def _add_random_ais(self) -> None:
        """adds random AIs to the population in order to reach the original population size"""
        missing_ais = self.population_size - len(self.population)
        ais = []
        for ai_number in range(len(self.ai_histories), len(self.ai_histories) + missing_ais):
            ai = AiPlayer(str(ai_number), Trainer._build_random_strategy(),
                          Trainer._build_random_strategy(), Trainer._build_random_strategy())
            # ai = AiPlayer(str(ai_number), SampleStrategies.bodo_quadratic_factor,
                          # SampleStrategies.bodo_linear_factor, SampleStrategies.bodo_bias)
            self.ai_histories[ai] = dict()
            self._add_event_to_ai_history(ai, "INITialization")
            ais.append(ai)
        self.population.extend(ais)
        assert (len(self.population) <= self.population_size)

    def _find_strongest_ai(self):
        """finds the highest scoring ai in a population"""
        max_points = float("-inf")
        strongest_ai = None
        for ai in self.population:
            if max_points < self.points_per_ai[ai]:
                max_points = self.points_per_ai[ai]
                strongest_ai = ai
        return strongest_ai, max_points

    def _compute_points_statistics(self):
        """calculates the average points that were scored in a generation"""
        sum_points = 0
        for ai in self.population:
            sum_points += self.points_per_ai[ai]
        avg_points = sum_points / self.population_size

        var_points = Trainer._compute_variance(self.points_per_ai.values(), avg_points)
        return avg_points, var_points

    def _compute_next_generation(self) -> None:
        """directs all steps that have to be done to create the next generation / a (partly) new population"""
        if self.generation > 0:
            self._select()
            self._recombine()
            self._mutate()
        self._add_random_ais()
        self._rank()  # &[(ai, list(reversed(list(self.ai_histories[ai].items())))) for ai in self.population]

    @staticmethod
    def _build_random_strategy():
        """builds any part of strategy (quadratic, linear, bias)"""
        return Trainer._adjust_strategy_range(np.random.rand(AiPlayer.strategy_length))

    @staticmethod
    def _adjust_strategy_range(array_or_value):
        width = Trainer.strategy_parameter_max - Trainer.strategy_parameter_min
        return array_or_value * width + Trainer.strategy_parameter_min

    def train(self, load_population_from_file, load_ai_as_population, save_population_in_file) -> None:
        """trains the AIs due to the parameters and returns the final and best AI"""
        if load_population_from_file:
            self._load_population()
        if load_ai_as_population:
            self._load_ai_as_population()
        else:
            assert (self.generation == 0 and self.population == [])

        best_ai = None
        max_points = float("-inf")
        stop_generation = self.num_generations + self.generation
        print(f"group_size = {self.group_size}, play_against_own_copies = {self.play_against_own_copies}, "
              f"population_size = {self.population_size}, survivor_rate = {self.survivor_rate}, \n"
              f"child_rate = {self.child_rate}, mutation_rate = {self.mutation_rate}, "
              f"mutation_copy_rate = {self.mutation_copy_rate}, lowest_variance_rate = {self.lowest_variance_rate}, "
              f"num_generations = {self.num_generations}")
        while self.generation < stop_generation:
            self._compute_next_generation()
            strong_ai, strong_ai_points = self._find_strongest_ai()
            if strong_ai_points > max_points:
                max_points = strong_ai_points
                best_ai = strong_ai
            avg_points, variance = self._compute_points_statistics()
            print(f"Generation: {self.generation} \t Max: {strong_ai_points:.0f} \t Avg: {avg_points:.0f} "
                  f"\t Var: {variance:.0f}")
            self.generation += 1
        avg_points, variance = self._play_against_own_copies(best_ai)
        print(f"Best AI \t\t\t\t\t Avg: {avg_points:.0f} \t Var: {variance:.0f}")
        print("Evolution finished")
        if save_population_in_file:
            self._save_final_population_and_best_ai(best_ai, avg_points)

    def _save_final_population_and_best_ai(self, best_ai, avg_points) -> None:
        """saves the final population and the best AI of a train cycle """
        pickle.dump((self.population, self.ai_histories, self.generation), open("final_population.dat", "wb"))
        pickle.dump((best_ai, self.ai_histories[best_ai]), open(f"best_ai_{avg_points:.0f}_points.dat", "wb"))

    def _load_ai_as_population(self):
        """loads a saved AI as a population, like AI loaded by SampleStrategies"""
        file = open("best_ai_67_points.dat", "rb")
        ai, ai_history = pickle.load(file)
        for _ in range(self.population_size):
            ai_copy = deepcopy(ai)
            self.population.append(ai_copy)
            self.ai_histories[ai_copy] = ai_history
        file.close()

    def _load_population(self) -> None:
        """loads the population that was saved"""
        file = open("final_population.dat", "rb")
        self.population, self.ai_histories, self.generation = pickle.load(file)
        file.close()

if __name__ == "__main__":
    trainer = Trainer()
    trainer.train(load_population_from_file=False, load_ai_as_population=False, save_population_in_file=True)


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

class Process():
    '''
        define a time series process
    '''

    def append(self, time_series):
        raise NotImplementedError(
            "method not defined for base class")

class PriceProcess(Process):

    def __init__(self, s=pd.Series()):
        self.s = s

    def append(self, s):
        '''
            extend two price process
        '''
        if len(self.s) == 0:
            self.s = self.s.append(s)
        else:
            adj_factor = self.s[-1] / 1
            s = s * adj_factor
            self.s = self.s.append(s)

    def __check_len(self, other):
        if (self.s.index != other.s.index).any() or \
            len(self.s.index) != len(other.s.index):
            raise ValueError(self.s.index,"!=",other.s.index)

    def __len__(self):
        return len(self.s)

    def plotvs(self, other, label=["portfolio", "index"],
        kind=["line", "line"], color=["blue", "orange"]):
        '''
            plot one price process against other process
        '''
        self.__check_len(other)

        plt.figure()
        self.s.plot(label=label[0], kind=kind[0], color=color[0])
        other.s.plot(label=label[1], kind=kind[1], color=color[1])
        plt.legend()
        plt.xlabel("Date")
        plt.ylabel("Net Value")
        plt.title("Tracking Performance")
        plt.show()


#!/usr/bin/python3
"""EXPORT TO JSON"""
import json
import requests
from sys import argv

if __name__ == '__main__':
    url = 'https://jsonplaceholder.typicode.com/'
    user_name = requests.get(url + f'users/{argv[1]}').json().get('username')
    todos = requests.get(url + f'todos?userId={argv[1]}').json()

    data = [
        {"task": f'{r.get("title")}',
         "completed": eval(f'{r.get("completed")}'),
         "username": user_name
         } for r in todos
         ]

    json_data = {argv[1]: [items for items in data]}
    with open(f'{argv[1]}.json', 'w', encoding='utf-8') as file:
        json.dump(json_data, file)