import scipy
import pyccl as ccl
import numpy as np
import pylab as plt
import itertools

import pymaster as nmt
from numpy import linalg
from sklearn.decomposition import PCA

class LinPkEmulator:
    def __init__(self, n_train, n_k, n_evecs, new = False):

        self.n_evecs = n_evecs
        self.k_arr = np.logspace(-4.7,5, n_k)
        self.wc_arr  = np.linspace(0.01, 1, 3*n_train)
        self.h_arr  = np.linspace(0.6,  0.9, n_train)

        if new:
            self.grid = []
            for i in range(len(self.wc_arr)):
                for j in range(len(self.h_arr)):
                    self.grid.append(self._get_theory_Pk(self.wc_arr[i], self.h_arr[j]))
            self.grid = np.array(self.grid)

            np.savetxt('Pk_grid.txt', self.grid)
        else:
            self.grid = np.loadtxt('Pk_grid.txt')

        self.log_grid = np.log(self.grid)
        self.mean_log_grid = np.mean(np.log(self.grid), axis=0) #avg over all cosmologies
        self.clean_grid = self.log_grid - self.mean_log_grid

        self.evecs = None
        self.w_arr = None
        self.w_emulator = None

        return

    def get_k_arr(self):
        return self.k_arr

    def get_evecs(self):
        if self.evecs is None:
            cov = np.cov(np.transpose(self.clean_grid))
            self.evals, self.evecs = linalg.eig(cov)
            #right to left vectors
            self.evecs = np.transpose(self.evecs)
            self.evecs = self.evecs[:self.n_evecs]
            #back to right
            self.evecs = np.transpose(self.evecs)
            #print(self.evals)
        return self.evecs

    def get_w_arr(self):
        if self.w_arr is None:
            self.w_arr = []
            l=0
            for i in range(len(self.wc_arr)):
                row = []
                for j in range(len(self.h_arr)):
                    #Pk_ij = self._get_theory_Pk(self.wc_arr[i], self.h_arr[j])
                    Pk_ij = self.grid[l]
                    l = l + 1
                    w_ij = self._get_weigths(Pk_ij)
                    row.append(w_ij)
                self.w_arr.append(row)

        self.w_arr = np.transpose(np.array(self.w_arr))
        return self.w_arr

    def get_w_emulator(self):
        if self.w_emulator is None:
            self.w_emulator={}
            self.get_w_arr()
            w =0
            for grid in self.w_arr:
                w_i_emulator = scipy.interpolate.interp2d( self.wc_arr, self.h_arr, grid, kind='cubic')
                self.w_emulator["w_{}".format(w)]= w_i_emulator
                w = w+1
        return self.w_emulator

    def _get_weigths(self, Pk):
        if self.evecs is None:
            self.get_evecs()
        y = np.log(Pk) - self.mean_log_grid
        return np.linalg.lstsq(self.evecs, y, rcond=None)[0]

    def _get_theory_Pk(self, wc, h):
        cosmo = ccl.Cosmology(Omega_c=wc, Omega_b=0.049, h=h, sigma8=0.81, n_s=0.96)
        return ccl.power.linear_matter_power(cosmo, self.k_arr, 1)

    def _get_evec_Pk(self, weigths , evecs, mean):
        return mean + np.sum(evecs*weigths, axis=1)

    def get_emulated_Pk(self, h, wc):
        #log lin_Pk
        w_vec = []
        self.get_w_emulator()
        for i in range(self.n_evecs):
            w_i=self.w_emulator["w_{}".format(i)](h, wc)[0]
            w_vec.append(w_i)
        w_vec = np.array(w_vec)
        Pk = self._get_evec_Pk(w_vec, self.evecs, self.mean_log_grid)
        return self.k_arr, np.array(Pk)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import itertools
import logging
from functools import lru_cache
from typing import Tuple

import numpy as np
from scipy.spatial import distance

from chord_eval.constants import TRIAD_REDUCTION
from chord_eval.data_types import ChordType, PitchType
from chord_eval.utils import (
    Distance,
    filter_noise,
    find_notes_matching,
    find_peaks,
    get_chord_pitches,
    get_dft_from_chord,
)

def SPS_distance(
    root1: int,
    root2: int,
    chord_type1: ChordType,
    chord_type2: ChordType,
    inversion1: int = 0,
    inversion2: int = 0,
    changes1: str = None,
    changes2: str = None,
    program1: int = 0,
    program2: int = 0,
    pitches: Tuple[int] = None,
    transform: str = "vqt",
    hop_length: int = 512,
    bins_per_octave: int = 60,
    n_mels: int = 512,
    noise_filtering: bool = False,
    peak_picking: bool = False,
) -> float:
    """
    Get the spectral pitch similarity (SPS) between two chords (composed of
    a root a ChordType and an inversion) using general MIDI programs.

    Parameters
    ----------
    root1 : int
        The root of the given first chord, as MIDI note number. If the chord
        is some inversion, the root pitch will be on this MIDI note, but there
        may be other pitches below it.

    root2 : int
        The root of the given second chord, as MIDI note number. If the chord
        is some inversion, the root pitch will be on this MIDI note, but there
        may be other pitches below it.

    chord_type1 : ChordType
        The chord type of the given first chord.

    chord_type2 : ChordType
        The chord type of the given second chord.

    inversion1 : int, optional
        The inversion of the first chord.
        The default is 0.

    inversion2 : int, optional
        The inversion of the second chord.
        The default is 0.

    changes1 : str
        Any alterations to the 1st chord's pitches, as a semi-colon separated string.
        Each alteration should be in the form "orig:new", where "orig" represents
        the original pitch that has been altered to "new". "orig" can also be blank
        for added pitches, and "new" can be prepended with "+" to indicate that a
        pitch occurs in an upper octave (e.g., a C7 chord with a 9th is represented
        by ":+1", using MIDI pitch). Note that TPC pitch does not allow for the
        representation of different octaves so any "+" is ignored.

    changes2 : str
        Any alterations to the 2nd chord's pitches.

    pitches : Tuple[int]
        A Tuple of possible absolute MIDI pitch numbers to use in each chord.
        If given, only a subset of these pitches will be included in the generated
        chord. Specifically, those which share a pitch class with any of the
        default chord tones. Note that this means that some default chord tones might
        not be present in the generated chord.

    program1 : int, optional
        The general MIDI program number used by the fluidsynth to synthesize
        the wave form of the MIDI data of the first chord (it uses the
        TimGM6mb.sf2 sound font file included with pretty_midi and synthesizes
        at fs=44100hz by default)
        The general MIDI program number (instrument index) is in [0, 127] :
        https://pjb.com.au/muscript/gm.html
        The default is 0.

    program2 : int, optional
        The general MIDI program number used by the fluidsynth to synthesize
        the wave form of the MIDI data of the second chord.
        The default is 0.

    transform : str, optional
        Indicates which transform to use to get the spectrogram of the
        synthesized instrument's notes. It can either be 'stft, 'cqt', 'vqt'
        or 'mel'/'melspectrogram'.
        The default is 'vqt'

    hop_length : int, optional
        Number of samples between successive CQT or VQT columns if transform is
        'cqt' or 'vqt'.
        the default is 512.

    bin_per_octave : int, optional
        Number of bins per octave if transform is 'cqt' or 'vqt'.
        the default is 60.

    n_mels : int, optional
        Number of Mel bands to generate if transform is 'mel' or 'melspectrogram'.
        the default is 512.

    noise_filtering : bool, optional
        If True, the function will filter out the noise of the two spectra
        using the filter_noise function.
        If peak_picking is True, it will automatically assign the True value
        to noise_filtering.
        The default is False.

    peak_picking : bool, optional
        If True, the function will isolate the peaks of each spectrum using
        the find_peaks function after filtering out the noise of each spectrum.
                If peak_picking is True, it will automatically assign the True value
        to noise_filtering.
        The default is False.

    Returns
    -------
    dist : float
        The cosine distance beween the spectra of the two synthesized chords
        (SPS) (in [0, 1]).
    """

    dft1 = get_dft_from_chord(
        root1,
        chord_type1,
        inversion=inversion1,
        changes=changes1,
        pitches=pitches,
        program=program1,
        transform=transform,
        hop_length=hop_length,
        bins_per_octave=bins_per_octave,
        n_mels=n_mels,
    )

    dft2 = get_dft_from_chord(
        root2,
        chord_type2,
        inversion=inversion2,
        changes=changes2,
        pitches=pitches,
        program=program2,
        transform=transform,
        hop_length=hop_length,
        bins_per_octave=bins_per_octave,
        n_mels=n_mels,
    )

    if peak_picking or noise_filtering:
        dft1 = [filter_noise(dft) for dft in dft1]
        dft2 = [filter_noise(dft) for dft in dft2]

        if peak_picking:
            dft1 = [find_peaks(dft) for dft in dft1]
            dft2 = [find_peaks(dft) for dft in dft2]

    sps = min([distance.cosine(d1, d2) for d1, d2 in itertools.product(dft1, dft2)])

    return sps

@lru_cache
def mechanical_distance(
    root1: int,
    root2: int,
    chord_type1: ChordType,
    chord_type2: ChordType,
    inversion1: int = 0,
    inversion2: int = 0,
    changes1: str = None,
    changes2: str = None,
    distance: Distance = Distance(),
    bass_distance: Distance = None,
    bass_weight: int = 1,
) -> float:
    """
    Get the voice leading distance between two chords : the number
    of semitones between the pitches of each chords.

    Parameters
    ----------
    root1 : int
        The root of the given first chord, as MIDI note number. If the chord
        is some inversion, the root pitch will be on this MIDI note, but there
        may be other pitches below it.

    root2 : int
        The root of the given second chord, as MIDI note number. If the chord
        is some inversion, the root pitch will be on this MIDI note, but there
        may be other pitches below it.

    chord_type1 : ChordType
        The chord type of the given first chord.

    chord_type2 : ChordType
        The chord type of the given second chord.

    inversion1 : int, optional
        The inversion of the first chord. The default is 0.

    inversion2 : int, optional
        The inversion of the second chord.. The default is 0.

    changes1 : str
        Any alterations to the 1st chord's pitches, as a semi-colon separated string.
        Each alteration should be in the form "orig:new", where "orig" represents
        the original pitch that has been altered to "new". "orig" can also be blank
        for added pitches, and "new" can be prepended with "+" to indicate that a
        pitch occurs in an upper octave (e.g., a C7 chord with a 9th is represented
        by ":+1", using MIDI pitch). Note that TPC pitch does not allow for the
        representation of different octaves so any "+" is ignored.

    changes2 : str
        Any alterations to the 2nd chord's pitches.

    distance : Distance
        The Distance measurement to use between notes. By default, this is
        semitone distance.

    bass_distance : Distance
        The Distance measurement to use between bass notes, if not the same as the
        distance parameter.

    bass_weight : int, optional
        The weight of the basses distance of the two chords. The default is 1.

    Returns
    -------
    total_steps : int
        The mechanical distance between the pitches of the two chords.
    """
    assert (
        distance.is_valid()
    ), "Given distance is not valid (some distances are negative)."
    if bass_distance is None:
        bass_distance = distance
    else:
        if not bass_distance.is_valid():
            logging.warning(
                "bass_distance is not valid (some values are negative). "
                "Defaulting to given distance."
            )
            bass_distance = distance

    # note number of the pitches in the chord
    notes1 = get_chord_pitches(
        root=root1,
        chord_type=chord_type1,
        pitch_type=PitchType.MIDI,
        inversion=inversion1,
        changes=changes1,
    )

    notes2 = get_chord_pitches(
        root=root2,
        chord_type=chord_type2,
        pitch_type=PitchType.MIDI,
        inversion=inversion2,
        changes=changes2,
    )

    # Bass distance
    bass_steps = bass_distance.distance_between(notes1[0], notes2[0]) * bass_weight

    # Other pitches
    upper_steps, match_ids = find_notes_matching(notes1, notes2, distance)

    # If the bass notes have been matched together, do not count again
    if (0, 0) in match_ids:
        upper_steps -= distance.distance_between(notes1[0], notes2[0])

    # Handle chords of different lengths
    if len(notes1) != len(notes2):
        big_chord, bc_idx, small_chord = (
            (notes1, 0, notes2) if len(notes1) > len(notes2) else (notes2, 1, notes1)
        )

        # Remove all pitches from the biggest chord that have already been matched
        big_chord = [
            big_chord[idx]
            for idx in range(len(big_chord))
            if idx not in set([pair[bc_idx] for pair in match_ids] + [0])
        ]

        # Find corresponding note in the smaller chord for each additional big chord note
        upper_steps += sum(
            [
                min(
                    [
                        distance.distance_between(bc_pitch, sc_pitch)
                        for sc_pitch in small_chord
                    ]
                )
                for bc_pitch in big_chord
            ]
        )

    return int(upper_steps + bass_steps)

@lru_cache
def tone_by_tone_distance(
    root1: int,
    root2: int,
    chord_type1: ChordType,
    chord_type2: ChordType,
    inversion1: int = 0,
    inversion2: int = 0,
    changes1: str = None,
    changes2: str = None,
    root_bonus: int = 1,
    bass_bonus: int = 1,
    pitch_type: PitchType = PitchType.MIDI,
) -> float:
    """
    Get the tone by tone distance between two chords: The average of the proporion
    of pitch classes in each chord that are unmatched in the other chord.

    For example, for C major and A minor triads, 1/3 of each one's pitch classes
    are unmatched, so the tone-by-tone distance is 1/3.

    For C major and C7, 0 of the C major's pitch classes are unmatched, and 1/4 of
    the C7's pitch classes are unmatched, so the distance is 1/8.

    Parameters
    ----------
    root1 : int
        The root of the given first chord, either as a MIDI note number or as
        a TPC note number (interval above C in fifths with F = -1, C = 0, G = 1, etc.).
        If the chord is some inversion, the root pitch will be on this pitch, but there
        may be other pitches below it.

    root2 : int
        The root of the 2nd chord.

    chord_type1 : ChordType
        The chord type of the given first chord.

    chord_type2 : ChordType
        The chord type of the given second chord.

    inversion1 : int, optional
        The inversion of the first chord.
        The default is 0.

    inversion2 : int, optional
        The inversion of the second chord.
        The default is 0.

    changes1 : str
        Any alterations to the 1st chord's pitches, as a semi-colon separated string.
        Each alteration should be in the form "orig:new", where "orig" represents
        the original pitch that has been altered to "new". "orig" can also be blank
        for added pitches, and "new" can be prepended with "+" to indicate that a
        pitch occurs in an upper octave (e.g., a C7 chord with a 9th is represented
        by ":+1", using MIDI pitch). Note that TPC pitch does not allow for the
        representation of different octaves so any "+" is ignored.

    changes2 : str
        Any alterations to the 2nd chord's pitches.

    root_bonus : int, optional
        Give additional accuracy to chords whose root notes match (and penalize those whose
        do not). 1 will give an additional bonus or penalty of equal weighting as each other
        note.

        For example, for a C major vs C minor comparison:
            - Default distance (root_bonus == 0): 1/3 (each chord matches 2/3 of their notes
              with the other chord).
            - With root_bonus == 1: 1/4 (each chord matches 2/3 of their notes with the other
              chord, plus an additional 1 "bonus match" up to 3/4 for their roots matching).

        For example, for C major vs Cmin7 comparison:
            - Default distance: 5 / 12 (0.417 -- the average of 1/3 from the C major triad and
              2/4 from the Cmin7).
            - With root_bonus == 1: 13/40 (0.325 -- the average of 1/4 from the C major triad
              and 2/5 from the Cmin7).

        For example, for C major vs A minor:
            - Default distance: 1/3
            - With root_bonus == 1: 1/2

    bass_bonus : int, optional
        Give additional accuracy to chords whose bass notes match. 1 will give an additional
        bonus of equal weighting as each other note.

        For example, for a C major vs C minor comparison:
            - Default distance (bass_bonus == 0): 1/3 (each chord matches 2/3 of their notes
              with the other chord).
            - With bass_bonus == 1: 1/4 (each chord matches 2/3 of their notes with the other
              chord, plus an additional 1 "bonus match" up to 3/4 for their bass notes matching).

        For example, for C major vs Cmin7 comparison:
            - Default distance: 5 / 12 (0.417 -- the average of 1/3 from the C major triad and
              2/4 from the Cmin7).
            - With bass_bonus == 1: 13/40 (0.325 -- the average of 1/4 from the C major triad
              and 2/5 from the Cmin7).

        For example, for C major vs A minor:
            - Default distance: 1/3
            - With bass_bonus == 1: 1/2

    pitch_type : PitchType
        The pitch type in which root notes are encoded, and using which tone-by-tone distance
        should be calculated. With MIDI (default), enharmonic equivalence is assumed. Otherwise,
        C# and Db count as different pitches.

    Returns
    -------
    distance : float
        The tone by tone distance between two chords.
    """

    def one_sided_tbt(
        note_set1: set,
        note_set2: set,
        root_matches: bool,
        bass_matches: bool,
        root_bonus: int,
        bass_bonus: int,
    ) -> float:
        """
        Get the one-sided tbt. That is, the proportion of chord 1's notes which are missing
        from chord2, including root and bass bonus.

        Parameters
        ----------
        note_set1 : set
            The set of pitch classes contained in chord 1.
        note_set2 : np.ndarray
            The set of pitch classes in chord 2.
        root_matches : bool
            True if the chord roots match. False otherwise.
        bass_matches : bool
            True if the chord bass notes match. False otherwise.
        root_bonus : int
            The root bonus to use (see the tone_by_tone_distance function's comments for details).
        bass_bonus : int
            The bass bonus to use (see the tone_by_tone_distance function's comments for details).

        Returns
        -------
        distance : float
            The one-sided tone-by-tone distance.
        """
        matches = len(note_set1.intersection(note_set2))

        if root_matches:
            matches += root_bonus
        if bass_matches:
            matches += bass_bonus

        return 1 - matches / (len(note_set1) + bass_bonus + root_bonus)

    notes1 = get_chord_pitches(
        root=root1,
        chord_type=chord_type1,
        pitch_type=PitchType.MIDI,
        inversion=inversion1,
        changes=changes1,
    )
    if pitch_type == PitchType.MIDI:
        notes1 %= 12

    notes2 = get_chord_pitches(
        root=root2,
        chord_type=chord_type2,
        pitch_type=PitchType.MIDI,
        inversion=inversion2,
        changes=changes2,
    )
    if pitch_type == PitchType.MIDI:
        notes2 %= 12

    root_matches = root1 == root2
    bass_matches = notes1[0] == notes2[0]

    distance = np.mean(
        [
            one_sided_tbt(
                set(notes1),
                set(notes2),
                root_matches,
                bass_matches,
                root_bonus,
                bass_bonus,
            ),
            one_sided_tbt(
                set(notes2),
                set(notes1),
                root_matches,
                bass_matches,
                root_bonus,
                bass_bonus,
            ),
        ]
    )

    return distance

def get_distance(
    distance: str,
    root1: int,
    root2: int,
    chord_type1: ChordType,
    chord_type2: ChordType,
    inversion1: int = 0,
    inversion2: int = 0,
    changes1: str = None,
    changes2: str = None,
    triad_reduction: bool = False,
    **kwargs,
) -> float:
    """
    Get the required distance between two chords : either the SPS, the voice leading,
    the tone by tone distance by calling the SPS_distance, mechanical_distance
    or the tone_by_tone_distance function respectively, or the binary distance.
    voice.

    Parameters
    ----------
    distance : str
        The name of the metric to use. It can be either 'SPS', 'mechanical',
        'tone by tone' or 'binary'.

    root1 : int
        The root pitch of the given first chord. With distance == "tone by tone",
        this may be either a MIDI pitch class or a TPC pitch class (see the
        tone_by_tone_distance function documentation for details). Otherwise,
        thei should be a MIDI pitch class.

    root2 : int
        The root of the given second chord.

    chord_type1 : ChordType
        The chord type of the given first chord.

    chord_type2 : ChordType
        The chord type of the given second chord.

    inversion1 : int, optional
        The inversion of the first chord.
        The default is 0.

    inversion2 : int, optional
        The inversion of the second chord.
        The default is 0.

    changes1 : str
        Any alterations to the 1st chord's pitches, as a semi-colon separated string.
        Each alteration should be in the form "orig:new", where "orig" represents
        the original pitch that has been altered to "new". "orig" can also be blank
        for added pitches, and "new" can be prepended with "+" to indicate that a
        pitch occurs in an upper octave (e.g., a C7 chord with a 9th is represented
        by ":+1", using MIDI pitch). Note that TPC pitch does not allow for the
        representation of different octaves so any "+" is ignored.

    changes2 : str
        Any alterations to the 2nd chord's pitches.

    triad_reduction : bool, optional
        If True, the function will reduce every input chords to their triad
        reduction.
        The default is False.

    **kwargs : TYPE
        Additional argument for the type of metric used.
        If distance is 'SPS', this will be arguments for the SPS_distance function.
        If distance is 'mechanical', this will be arguments for the mechanical_distance function.
        If distance is 'tone_by_tone', this will be arguments for the
        tone_by_tone_distance function.

    Raises
    ------
    ValueError
        if distance is something other than 'SPS', 'mechanical', 'tone by tone'
                or 'binary'.

    Returns
    -------
    float
        The corresponding distance between the two chords.

    """
    if triad_reduction:
        chord_type1 = TRIAD_REDUCTION[chord_type1]
        chord_type2 = TRIAD_REDUCTION[chord_type2]

    if distance == "SPS":
        return SPS_distance(
            root1=root1,
            root2=root2,
            chord_type1=chord_type1,
            chord_type2=chord_type2,
            inversion1=inversion1,
            inversion2=inversion2,
            changes1=changes1,
            changes2=changes2,
            **kwargs,
        )

    elif distance == "mechanical":
        return mechanical_distance(
            root1=root1,
            root2=root2,
            chord_type1=chord_type1,
            chord_type2=chord_type2,
            inversion1=inversion1,
            inversion2=inversion2,
            changes1=changes1,
            changes2=changes2,
            **kwargs,
        )

    elif distance == "tone by tone":
        return tone_by_tone_distance(
            root1=root1,
            root2=root2,
            chord_type1=chord_type1,
            chord_type2=chord_type2,
            inversion1=inversion1,
            inversion2=inversion2,
            changes1=changes1,
            changes2=changes2,
            **kwargs,
        )

    elif distance == "binary":
        return (
            0
            if root1 == root2
            and chord_type1 == chord_type2
            and inversion1 == inversion2
            and changes1 == changes2
            else 1
        )

    else:
        raise ValueError(
            "distance must be 'SPS', 'mechanical', 'tone by tone' or 'binary'."
        )


'''
Created on 31 de dez de 2020

@author: klaus
'''
import os
import re

import numpy as np
import os.path as osp
import pandas as pd
from tf_labelprop.experiment.prefixes import AFFMAT_PREFIX, NOISE_PREFIX, ALG_PREFIX, FILTER_PREFIX, GENERAL_PREFIX, OUTPUT_PREFIX, INPUT_PREFIX
from tf_labelprop.logging.logger import log as LOG
from tf_labelprop.output.folders import CSV_FOLDER, PLOT_FOLDER, RESULTS_FOLDER, \
    VIS_FOLDER

def fix_underscore(x):
    return str(x).replace('_','\\textunderscore ')

def default_translate():
    dct = {
        f'{INPUT_PREFIX}dataset': 'Dataset',
        f'{INPUT_PREFIX}labeled_percent': '\\% labels',
        f'{INPUT_PREFIX}num_labeled': 'labels',
        f'{INPUT_PREFIX}benchmark': 'Benchmark',

        f'{AFFMAT_PREFIX}dist_func': 'Affmat dist.',
        f'{AFFMAT_PREFIX}sigma': '$\\sigma$',
        f'{AFFMAT_PREFIX}k': 'k',
        f'{AFFMAT_PREFIX}row_normalize':'Normalize rows',

        f'{FILTER_PREFIX}mu':'$\\mu$',
        f'{FILTER_PREFIX}filter':'Filter',

        f'{NOISE_PREFIX}corruption_level':'Noise \\%',
        f'{NOISE_PREFIX}deterministic':'Deterministic noise',
        f'{NOISE_PREFIX}type':'Noise type',

        f'{ALG_PREFIX}algorithm':'Classifier',
        f'{ALG_PREFIX}alpha':'$\\alpha$',
        f'{ALG_PREFIX}mu':'$\\mu$',
        f'{ALG_PREFIX}optimize_labels':'opt.l.',
        f'{ALG_PREFIX}custom_conv':'c.conv.',

        f'{OUTPUT_PREFIX}acc':'Acc.',
        f'{OUTPUT_PREFIX}acc_labeled':'L.Acc.',
        f'{OUTPUT_PREFIX}acc_unlabeled':'Ul.Acc.',

        f'{OUTPUT_PREFIX}CMN_acc':'Acc. (w/ CMN)',
        f'{OUTPUT_PREFIX}CMN_acc':'Acc. (labeled, w/ CMN)',
        f'{OUTPUT_PREFIX}CMN_acc':'Acc. (unlabeled, w/ CMN)',

        f'{OUTPUT_PREFIX}affmat_time':'Time (affmat)',
        f'{OUTPUT_PREFIX}classifier_time':'Time (classifier)',
        f'{OUTPUT_PREFIX}filter_time':'Time (classifier)',
        f'{OUTPUT_PREFIX}noise_time':'Time (noise)'
        }
    dct2 = dict(dct)
    for k in dct.keys():
        if k.startswith(OUTPUT_PREFIX):
            for x,y in zip(['mean','median','min','max','sd'],
                           ['','median','min','max','sd']):
                if len(y) > 0:
                    dct2[k+f'_{x}'] = dct[k] + f'({y})'
                else:
                    dct2[k+f'_{x}'] = dct[k]
    return dct2

def tex_table(df,y_vars,t_vars,vdiv_vars,hdiv_vars,c_vars,use_std,OUT_FOLDER,dct_translate = default_translate(),fixed_vars=[],bold_max=True):
    """
        Creates and saves a table in LaTeX corresponding to the observation of aggregated variables.

            Args:
                df (pd.DataFrame): Dataframe containing aggregated observations
                y_vars (List[str]): List of output values to be present.
                t_vars (List[str]): Variables used to define tables. There'll be one table corresponding to each combination of these variables.
                vdiv_vars (List[str]):  List of variables that correspond to vertical divisions.
                hdiv_vars (List[str]): Variables used to represent horizontal rules. In each table, there'll be a horizontal rule separating each combination of these variables.
                c_vars (List[str]): Variables used to represent configurations. In each horizontal rule defined by `hdiv_vars`, there'll be one line for each combination of `c_vars` variables.
                use_sd (bool): If ``True``, will show +- 1 standard deviation from the mean as a shaded region.
                OUT_FOLDER (str): Folder to which the visualization will be saved.
                dct_translate (Dict[str,str]): Maps the columns of the DataFrame to strings that are more readable.
                fixed_vars (List[str]): List of fixed variables
                bold_max (bool): If ``True``, automatically makes the text of the largest value

    """

    if len(t_vars) > 0:
        for g, g_df in df.groupby(t_vars):
            if len(t_vars) == 1:
                g = [g]
            else:
                g = list(g)

            _path = [OUT_FOLDER] + [x+'='+y for x,y in sorted(zip(t_vars,g))]
            _path = osp.join(*_path)
            tex_table(g_df,y_vars,[],vdiv_vars,hdiv_vars,c_vars,use_std,
                      OUT_FOLDER=_path,
                      dct_translate=dct_translate,fixed_vars=t_vars+fixed_vars,bold_max=bold_max)
    else:
        """
            Base case
        """
        fixed_vars_desc =  ';\n'.join([c + ":   "+  ';'.join(pd.unique(df[c]).astype(np.str)) for c in fixed_vars ])
        _tex = []

        """
            Title
        """
        print(OUT_FOLDER)
        if not osp.isdir(OUT_FOLDER):
            os.makedirs(OUT_FOLDER)
        with open(osp.join(OUT_FOLDER,'table.txt'), "w") as _tex:
            _tex.write("\\begin{table}[h] \\tiny")

            df[:] = df[:].fillna('---')
            for c in hdiv_vars+c_vars+vdiv_vars:
                if df[c].dtype.kind in 'biufc':
                    df[c] = df[c].fillna(-1)
                else:
                    df[c] = df[c].fillna('NA').astype('string')
                df[c] = pd.Categorical(df[c]).as_ordered()
            df[:] = df.sort_values(by=hdiv_vars+c_vars+vdiv_vars)

            """
                -------------------
                    Create combination of all vdiv_vars with y_vars
                ----------------
            """
            lvls= [list(np.sort(pd.unique(df[y]))) for y in vdiv_vars]
            lvls.append(y_vars)
            import itertools
            comb_v = list(itertools.product(*lvls))
            comb_v = [list(zip(vdiv_vars + [OUTPUT_PREFIX],x)) for x in comb_v]
            comb_v_dct = [dict(x) for x in comb_v]

            _tex.write("\\begin{tabular}{@{}|" +\
                        f'{"c"* len(c_vars+hdiv_vars)}' +
                        f"{'|c'* len(comb_v)}" + "|@{}} " + "\cline{" +\
                        f"{len(c_vars+hdiv_vars)+1}-{len(c_vars+hdiv_vars) + len(comb_v)}" + "} \n")

            """ Create Table HEADER """
            rep = 1
            multicolumn_len = len(comb_v)
            for v in vdiv_vars:
                L = ["\multicolumn{" + str(len(c_vars+hdiv_vars))+"}{c|}{}"]
                unique = pd.unique(df[v])
                multicolumn_len = multicolumn_len // (len(unique))
                print(v)
                print(multicolumn_len)
                for _ in range(rep):
                    for u in unique:
                        L.append('\\multicolumn{'+str(multicolumn_len) + "}{c|}{"+ f"{dct_translate.get(v,v)}={u}"+"}")
                _tex.write(fix_underscore(' & '.join(L)))
                _tex.write('\\\\ \n')
                rep *= len(unique)
            #_tex.write(fix_underscore('&'.join(c_vars+hdiv_vars+comb_v_cols)))
            _tex.write("\\hline \n")
            L = [dct_translate.get(c,c) for c in hdiv_vars+c_vars] + ([dct_translate.get(y,y) for y in pd.unique(y_vars)]\
                                                                 * ( len(comb_v)//pd.unique(y_vars).shape[0]    ))
            _tex.write(fix_underscore(' & '.join(L)))
            _tex.write("\\\\ \\hline  \n")
            def groupby(df,vrs):
                if len(vrs) == 1:
                    for v, v_df in df.groupby(vrs):
                        yield [v], v_df
                elif len(vrs) == 0:
                    yield [], df
                else:
                    for v, v_df in df.groupby(vrs):
                        yield v, v_df

            df = df.sort_values(hdiv_vars+c_vars+vdiv_vars)
            N_row = df.shape[0]
            row_counter = 0
            for h, h_df in groupby(df,hdiv_vars):
                oldh_df = h_df

                def _apply(h_df,out_var):
                    def is_number(x):
                        try:
                            x = float(x)
                        except:
                            return False
                        return True
                    def as_number(x):
                        try:
                            x = float(x)
                        except:
                            return 0.0
                        return x
                    b = np.where([is_number(x) for x in h_df[out_var]])[0]
                    if b.shape[0] > 0:
                        maxval = np.argmax([as_number(x) for x in h_df[out_var]],axis=0)
                    else:
                        maxval = -1
                    y = out_var
                    if isinstance(out_var,str):
                        out_var = [out_var]
                    def check_number(x):
                        try:
                            x = np.round(100*float(x),2)
                        except:
                            x = x
                            #print(x)
                        return x
                    out_var_r = [check_number(x) for x in h_df[out_var].values]

                    h_df[out_var] = out_var_r
                    if y.endswith('_mean'):
                        old_y = y
                        y = y[:-len('_mean')]
                        y = y+'_sd'
                        #print(list(zip(h_df[old_y].values,h_df.loc[:,y].values)))

                        def f(x):
                            return "$\\mathbf{"+ str(x[0]) + " \\pm " + str(check_number(x[1])) + "}$"
                        def g(x):
                            return "$" + str(x[0]) + " \\pm " + str(check_number(x[1]))+ "$"
                        print("==================")
                        print(h_df[old_y].values)
                        print(h_df[y].values)
                        print([f(x) if j == maxval   else f(x)  for j,x in enumerate(zip(h_df[old_y].values,h_df[y].values))])
                        h_df[old_y] = [f(x) if j == maxval   else g(x)  for j,x in enumerate(zip(h_df[old_y].values,h_df[y].values))]

                    else:
                        h_df[y] = ["$\\mathbf{"+ x +"}$" if j == maxval   else '$'+x+'$' for j,x in enumerate(h_df[y].values)]
                    return h_df
                for y in y_vars:
                    from functools import partial
                    h_df[y] = h_df.groupby(vdiv_vars).apply(partial(_apply,out_var=y))[y]

                for c, c_df in groupby(h_df,c_vars):
                    if c_df.shape[0] > 0:
                        row = [str(y)  for y in h] + [str(y)  for y in c]
                        #print("==========================")
                        for v, v_df in groupby(c_df,vdiv_vars):
                            print(v_df[y])
                            _r = [str(v_df[y].values[0])  if v_df[y].shape[0] > 0 else "---" for y in y_vars ]
                            row.extend(_r)
                        _tex.write(' & '.join(row))
                        #print(row)
                        row_counter += c_df.shape[0]
                        if row_counter < N_row:
                            _tex.write("\\\\ \n")

                if row_counter < N_row:
                    _tex.write("\\hline \n")

            """
                FINALIZE Table
            """
            _tex.write("\\end{tabular}\n ")

            _tex.write("\n \\caption{" + fix_underscore(fixed_vars_desc) +  "}\n")
            _tex.write("\\end{table}")

def line_plot(df,x_var,y_var,agg_var,p_vars,l_vars, use_std,OUT_FOLDER,dct_translate = default_translate(),fixed_var_desc=''):
    """
        Creates and saves a lineplot.
            Args:
                df (pd.DataFrame): Dataframe containing aggregated observations
                x_var (str): Variable to be plotted on x axis.
                y_var (str): Variable to be plotted on y axis.
                agg_var (str): Which aggregation to use. One of ``mean``,``median``,``min``,``max``.
                p_vars (List[str]): Variables used to define plots. There'll be one plot corresponding to each combination of these variables.
                l_vars (List[str]): Variables used to represent levels. When the variables in `p_vars` are fixed, there'll be one line for each combination of `l_vars` variables.
                use_sd (bool): If ``True``, will show +- 1 standard deviation from the mean as a shaded region.
                OUT_FOLDER (str): Folder to which the visualization will be saved.
                dct_translate (Dict[str,str]): Maps the columns of the DataFrame to strings that are more readable.
                fixed_var_desc (str): A string containing the description of the fixed variables of the dataframe. Will be placed at the bottom of the plot.

    """

    if len(p_vars) > 0:
        for g, g_df in df.groupby(p_vars):
            if len(p_vars) == 1:
                g = [g]
            else:
                g = list(g)

            _path = [OUT_FOLDER] + [x+'='+y for x,y in sorted(zip(p_vars,g))]
            _path = osp.join(*_path)
            line_plot(g_df,x_var,y_var,agg_var,[],l_vars,use_std,OUT_FOLDER=_path,dct_translate=dct_translate,fixed_var_desc=fixed_var_desc)
    else:
        """
            Base case
        """
        import matplotlib.pyplot as plt
        import matplotlib
        #matplotlib.use('TkAgg')

        if len(l_vars) == 0:
            df[' '] = ' '
            l_vars = [' ']

        """"
        ---------------------------------------
            DEFINE FONT SIZE
        --------------------------------

        """
        SMALL_SIZE = 14
        MEDIUM_SIZE = 20
        BIGGER_SIZE = 22
        BIGGEST_SIZE = 28

        plt.rc('font', size=BIGGER_SIZE)          # controls default text sizes
        plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
        plt.rc('axes', labelsize=BIGGEST_SIZE)    # fontsize of the x and y labels
        plt.rc('xtick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
        plt.rc('ytick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
        plt.rc('legend', fontsize=BIGGER_SIZE)    # legend fontsize
        plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title
        plt.rcParams.update({
            "text.usetex": True,
        })

        import re
        number_newlines = len(re.findall('\\n',fixed_var_desc))
        fig, ax = plt.subplots(1,figsize=(20,10+number_newlines//2))
        df[l_vars] = df[l_vars].fillna('---')
        for l, l_df in df.groupby(l_vars):
            if len(l_vars) == 1:
                l = [l]
            else:
                l = list(l)

            if len(l_df[x_var].dropna()) == 0:
                l_df[x_var] = '-'
            elif isinstance(l_df[x_var].dropna().values[0],str):
                l_df[x_var] = l_df[x_var].fillna('-')

            l_inst = ';'.join([dct_translate[x]+'='+fix_underscore(y) for x,y in sorted(zip(l_vars,l))])
            print(l_df[x_var].values)
            print(l_df[f'{y_var}_{agg_var}'].values)
            print(fix_underscore(l_inst))
            p = ax.plot(l_df[x_var].values,l_df[f'{y_var}_{agg_var}'].values,marker='o',
                    label=fix_underscore(l_inst),
                     linewidth=4, markersize=12)
            facecolor = matplotlib.colors.to_rgba(p[0].get_color(),alpha=0.1)
            edgecolor = matplotlib.colors.to_rgba(p[0].get_color(),alpha=0.8)

            if use_std:
                ax.fill_between(l_df[x_var],
                                l_df[f'{y_var}_mean'] - l_df[f'{y_var}_sd'],
                                l_df[f'{y_var}_mean'] + l_df[f'{y_var}_sd'],
                                facecolor = facecolor,
                                edgecolor= edgecolor
                                )

        print(OUT_FOLDER)
        plot_subtitle = fix_underscore(fixed_var_desc)
        ax.set_xlabel(fix_underscore(dct_translate.get(x_var,x_var)) +\
                      '\n\n' + plot_subtitle)
        ax.set_ylabel(fix_underscore(dct_translate.get(f'{y_var}_{agg_var}',f'{y_var}_{agg_var}')) )
        fig.tight_layout()
        plt.grid()

        plt.legend()
        if not osp.isdir(OUT_FOLDER):
            os.makedirs(OUT_FOLDER)
        fig.savefig(osp.join(OUT_FOLDER,f'PLOT_{y_var}__from__{x_var}.png'))

def interactive_line_plot(df,OUT_FOLDER):

    print("Dataframe columns:")
    print('\n'.join(list(df.columns)))

    _y_var = None
    while _y_var is None:
        try:
            _y_var = input(f"Type dependent output variable (e.g. {OUTPUT_PREFIX}acc)")
            assert _y_var + '_mean' in df.columns
        except:
            _y_var = None

    _agg_var = None
    while _agg_var is None:
        try:
            _agg_var = input(f"Type dependent output aggregation variable (mean,median,min,max)")
            assert _agg_var in ['mean','median','min','max']
        except:
            _agg_var = None

    _show_std = None
    while _show_std is None:
        try:
            _show_std = input(f"Show standard deviation (y/n)")
            assert _show_std in ['y','n']
            _show_std = True if _show_std == 'y' else False
        except:
            _show_std = None

    _x_var = None
    while _x_var is None:
        try:
            _x_var = input(f"Type INdependent variable (e.g. {NOISE_PREFIX}corruption_level)")
            assert _x_var in df.columns
        except:
            _x_var = None

    print("=========================================================================================================")
    print("The following variables have more than one value:")
    _changed_vars = [c for c in df.columns if (not c.startswith(OUTPUT_PREFIX)) and not (c in [_x_var,_y_var])]
    _changed_vars = [c for c in _changed_vars if len(pd.unique(df[c]))>1 ]

    _changed_desc = [c + ":   "+  ';'.join(pd.unique(df[c]).astype(np.str)) for c in _changed_vars ]
    print('\n'.join(_changed_desc))
    print("=========================================================================================================")
    print("For each of these variables, you'll have to select one of: ")
    print("l - use as level")
    print("g - use as group (there'll by one plot for each combination of such variables)")
    print("i - ignore (replaces all occurences with first value)")

    l_vars = []
    g_vars = []
    for v in _changed_vars:
        ch = 0
        while ch == 0:
            try:
                ch = input(f"Enter the variable type for {v}\n")
                assert ch in ['l','g','i']
                if ch == 'l':
                    l_vars.append(v)
                elif ch == 'g':
                    g_vars.append(v)
                else:
                    df[v] = pd.unique(df[v])[0]
            except:
                ch=0
    fixed_vars = list(set([c for c in df.columns if (not c.startswith(OUTPUT_PREFIX)) and not (c in [_x_var,_y_var])]).difference(g_vars+l_vars))
    fixed_vars_desc =  '\n'.join([c + ":   "+  ';'.join(pd.unique(df[c]).astype(np.str)) for c in fixed_vars ])
    line_plot(df, _x_var,_y_var,_agg_var,g_vars, l_vars, _show_std, OUT_FOLDER,fixed_var_desc=fixed_vars_desc)

def interactive_tex_table(df,OUT_FOLDER):

    print("Dataframe columns:")
    print('\n'.join(list(df.columns)))

    _out_var = None
    while _out_var is None:
        try:
            _out_var = input(f"Type comma-separated output variables w/ aggregation (e.g. {OUTPUT_PREFIX}acc_mean)")
            _out_var = _out_var.split(',')
            for y in _out_var:
                assert y in df.columns
                assert y.startswith(OUTPUT_PREFIX)
        except:
            _out_var = None

    _show_std = None
    while _show_std is None:
        try:
            _show_std = input(f"Use standard deviation  for *_mean output variables (y/n)")
            assert _show_std in ['y','n']
            _show_std = True if _show_std == 'y' else False
        except:
            _show_std = None

    _x_var = None
    while _x_var is None:
        try:
            _x_var = input(f"Type comma-separated input variables corresponding to VERTICAL divisions variable (e.g. {INPUT_PREFIX}labeled_percent)")
            _x_var = _x_var.split(',')
            for x in _x_var:
                assert x in df.columns
                assert not (x.startswith(OUTPUT_PREFIX))
        except:
            _x_var = None

    print("=========================================================================================================")
    print("The following variables have more than one value:")
    _changed_vars = [c for c in df.columns if (not c.startswith(OUTPUT_PREFIX)) and not (c in _x_var+_out_var)]
    _changed_vars = [c for c in _changed_vars if len(pd.unique(df[c]))>1 ]

    _changed_desc = [c + ":   "+  ';'.join(pd.unique(df[c]).astype(np.str)) for c in _changed_vars ]
    print('\n'.join(_changed_desc))
    print("=========================================================================================================")
    print("For each of these variables, you'll have to select one of: ")
    print("h - variable divides the table with horizontal rules")
    print("c - variable does not divide the table with horizontal rule")

    print("g - use as group (there'll by one plot for each combination of such variables)")
    print("i - ignore (replaces all occurences with first value)")

    c_vars = []
    h_vars = []
    g_vars = []
    for v in _changed_vars:
        ch = 0
        while ch == 0:
            try:
                ch = input(f"Enter the variable type for {v}\n")
                assert ch in ['c','h','g','i']
                if ch == 'c':
                    c_vars.append(v)
                elif ch == 'g':
                    g_vars.append(v)
                elif ch == 'h':
                    h_vars.append(v)
                else:
                    df[v] = pd.unique(df[v])[0]
            except:
                ch=0
    fixed_vars = list(set([c for c in df.columns if (not c.startswith(OUTPUT_PREFIX)) and not (c in _x_var+_out_var)]).difference(g_vars+c_vars+h_vars))

    tex_table(df=df,
              y_vars=_out_var,
              t_vars=g_vars,
              vdiv_vars=_x_var,
              hdiv_vars=h_vars,
              c_vars=c_vars,
             use_std=_show_std,
             OUT_FOLDER=OUT_FOLDER,
            fixed_vars =fixed_vars,
            bold_max=True)

def create_visualization(FOLDER=CSV_FOLDER,OUT_FOLDER=VIS_FOLDER):
    j_str = '_joined.csv'
    _len = len(j_str)
    fs = [x[:-_len] for x in os.listdir(FOLDER) if x.endswith(j_str) ]
    fs = np.sort(fs)

    """
        -------------------------------------------
            CHOOSE CSVs
        -----------------------------------
    """
    ch = []
    while len(ch) == 0:
        try:
            for i,f in enumerate(fs):
                print(f"{i} - {f}")
            ch = input("Enter comma-separated numbers corresponding to desired experiment\n")
            ch = str.split(ch,',')
            for x in ch:
                assert str.isnumeric(x)
            ch = [osp.join(FOLDER,fs[int(x)]+j_str) for x in ch]
        except:
            ch = []
    print("Selected files:\n {}".format('\n'.join(ch)))
    df = pd.concat([pd.read_csv(x) for x in ch],axis=0,ignore_index=True)

    if 'Unnamed: 0' in df.columns:
        df = df.drop(['Unnamed: 0'],axis=1)
    df = df.drop([c for c in df.columns if c.endswith("_values")],axis=1)

    """
    ----------------------------------------
        CHOOSE Output folder
    -------------------------------------

    """

    _subf = None
    while _subf is None:
        try:
            _subf = input(f"Type subfolder name to place visualizations within f{OUT_FOLDER}:\n")
            OUT_FOLDER = osp.join(OUT_FOLDER,_subf)
        except:
            _subf= None

    print(f"Will use save visualizations to {OUT_FOLDER}.")

    """
        -------------------------------------------
            CHOOSE Visualization
        -----------------------------------
    """
    ch = 0
    while ch == 0:
        try:
            ch = (input("Which visualization?  1-Line plot; 2-LaTeX Table (Use comma-separated values to use multiple)"))
            ch = [int(x) for x in ch.split(',')]
            for x in ch:
                assert x in [1,2]
        except:
            ch = 0
    print(ch)
    if 1 in ch:
        print("Line Plot Creation:\n")
        interactive_line_plot(df,OUT_FOLDER)

    if 2 in ch:
        print("LaTeX Table Creation:\n")
        interactive_tex_table(df, OUT_FOLDER)

if __name__ == "__main__":
    from io import StringIO
    import sys
    """
    s = StringIO('\n'.join(['0,1,2,3,4','exp_siis_unlabeled','2','out_acc_mean',
                            'y',
                            'noise_corruption_level']+list('igiciichiccci')))
    sys.stdin = s
    """
    create_visualization()
    sys.stdin = sys.__stdin__


#함수 만들
def add(a,b):
    c = a+b
    d = a - b
    return c,d

print(add(3,2))

#약수만들기 - 반복문을 돌다가 나눠지는 값은 약수가 있기에 소수가 아님 return False를 돌려주고 함수 종료
def isPrime(x):
    #True and False return
    for i in range(2, x):
        if x % i == 0:
            return False
    return True

a = [12,13,7,9,19]

for x in a:
    #x가 소수이면 print
    if isPrime(x):
        print(x, end=' ')

1. 기초 문법/11. 함수만들기.py


from sys import stdin, stdout
from collections import deque

read = stdin.readline
write = stdout.write

n = int(read().rstrip())
data = list(map(int, read().rsplit(' ')))
ans = [10000000000] * n

ans[0] = 0
for i in range(n):
    for j in range(1, data[i]+1):
        if i + j >= n: continue
        if ans[i + j] > ans[i] + 1:
            ans[i + j] = ans[i] + 1

if ans[-1] == 10000000000:
    ans[-1] = -1
write(str(ans[-1]))


from __future__ import print_function
import MDAnalysis
import numpy as np
from newanalysis.diffusion import msdCOM, unfoldTraj, msdMJ
import time

###########################################################################################################################
# Trajectory
###########################################################################################################################
base='../../data/emim_dca_equilibrium/'
psf=base+'emim_dca.psf'
#Check PSF:
if np.array_equal(MDAnalysis.Universe(psf).atoms.masses, MDAnalysis.Universe(psf).atoms.masses.astype(bool)):
    print("Used wrong PSF format for S0 state (masses unreadable!)")
    sys.exit()

skip=20
u=MDAnalysis.Universe(psf,"../../data/emim_dca_equilibrium/emim_dca.dcd")
boxl = np.float64(round(u.coord.dimensions[0],4))
dt=round(u.trajectory.dt,4)

n = int(u.trajectory.n_frames/skip)
if u.trajectory.n_frames%skip != 0:
    n+=1

###########################################################################################################################
# molecular species
###########################################################################################################################
sel_cat1  = u.select_atoms("resname EMIM")
ncat1     = sel_cat1.n_residues
com_cat1  = np.zeros((n,ncat1,3),dtype=np.float64)
print("Number EMIM   = ",ncat1)

sel_an1  = u.select_atoms("resname DCA")
nan1     = sel_an1.n_residues
com_an1  = np.zeros((n,nan1, 3),dtype=np.float64)
print("Number DCA   = ",nan1)

###########################################################################################################################
# Analysis
###########################################################################################################################
ctr=0
start=time.time()
print("")

for ts in u.trajectory[::skip]:
    print("\033[1AFrame %d of %d" % (ts.frame,len(u.trajectory)), "\tElapsed time: %.2f hours" % ((time.time()-start)/3600))
    com_an1[ctr] = sel_an1.center_of_mass(compound='residues')
    com_cat1[ctr]= sel_cat1.center_of_mass(compound='residues')
    ctr+=1

print("unfolding coordinates ..")
unfoldTraj(com_cat1,boxl)
unfoldTraj(com_an1, boxl)

print("calculating msd ..")
msd_an1  = msdCOM(com_an1)
msd_cat1 = msdCOM(com_cat1)

fan1=open("msd_an.dat",'w')
fcat1=open("msd_cat.dat",'w')

for i in range(ctr):
    fan1.write("%5.5f\t%5.5f\n"  % (i*skip*dt, msd_an1[i]))
    fcat1.write("%5.5f\t%5.5f\n" % (i*skip*dt, msd_cat1[i]))

print("calculating msdMJ ..")
msdmj = msdMJ(com_cat1,com_an1)
f1=open("msdMJ.dat",'w')
for i in range(ctr):
    f1.write("%5.5f\t%5.5f\n" % (i*skip*dt, msdmj[i]))
f1.close()


"""
Module containing functions pertaining to training models on watch data.

Author: Nate Burley
"""

import numpy as np
import pandas as pd
import pickle
import datetime
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score

def train_extra_trees(clf: ExtraTreesClassifier, X_train: np.array, y_train: np.array, save: bool=True, grid_search:bool=True) -> ExtraTreesClassifier:
    """
    Function to train an optimal Extra Trees Classifier using a grid search.
    Helpful example: https://www.kaggle.com/eikedehling/extra-trees-tuning
    More info on hyperparameters here: https://machinelearningmastery.com/extra-trees-ensemble-with-python/
    """

    if grid_search:
        # Define our grid search
        gsc = GridSearchCV(
            estimator=clf,
            param_grid={
                'n_estimators': range(50, 150, 50),
                'max_features': ['sqrt', 'log2', None],
                # 'min_samples_leaf': range(1,51,10),
                # 'min_samples_split': range(2,53,10)
            },
            scoring='accuracy',
            verbose=2
        )

        # Train the classifier and find the optimal parameters
        grid_result = gsc.fit(X_train, y_train)

        # Display best parameters and score
        print("Best: %f using %s" % (grid_result.best_score_, grid_result.best_params_))

        # Display the results of training across all parameters
        for test_mean, param in zip(
                grid_result.cv_results_['mean_test_score'],
                # grid_result.cv_results_['mean_train_score'],
                grid_result.cv_results_['params']):
            print("Test : %f with: %r" % (test_mean, param))

        # Construct and fit optimal model
        model = ExtraTreesClassifier(**grid_result.best_params_)
        model.fit(X_train, y_train)

        # Save the optimal classifier!
        if save:
            outfile = open('saved_models/best_extra_trees.pickle', 'wb')
            pickle.dump(model, outfile)

        # Return the optimal model
        return model

    else:
        # Construct and fit optimal model
        model = ExtraTreesClassifier()
        model.fit(X_train, y_train)

        # Save the classifier!
        if save:
            outfile = open('saved_models/default_extra_trees.pickle', 'wb')
            pickle.dump(model, outfile)

        # Return the optimal model
        return model

def evaluate_model(model: ExtraTreesClassifier, X_test: np.array, y_test: np.array, save=False) -> None:
    # Get our model's predictions on the test data
    y_test_pred = model.predict(X_test)

    # Generate a classification report
    print("CLASSIFICATION REPORT")
    report = classification_report(y_true=y_test, y_pred=y_test_pred, output_dict=True)
    report = pd.DataFrame(report).transpose()
    print(report.head(10))
    if save:
        report.to_excel('extratrees_classification_report.xlsx')

    # Print confusion matrix
    print("\nCONFUSION MATRIX")
    cm = confusion_matrix(y_true=y_test, y_pred=y_test_pred)
    print(cm)

    # And print the global accuracy
    print("\nACCURACY SCORES")
    print(accuracy_score(y_true=y_test, y_pred=y_test_pred))


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

import json
from uuid import uuid1
from datetime import date, time, datetime, timedelta, timezone
import icalendar

TERMWEEKS = 16

with open("class.json", 'r') as f:
    entry = json.load(f)

calendar = icalendar.Calendar()
calendar['version'] = '2.0'
calendar['prodid'] = '-//THU//Syllabus//CN'

for klass in entry:
    event = icalendar.Event()

    kl_start = 0
    kl_interval = 1
    if(klass['ClassWeeks'] == '全'):
        kl_count = TERMWEEKS
    elif(klass['ClassWeeks'] == '前八'):
        kl_count = TERMWEEKS // 2
    elif(klass['ClassWeeks'] == '后八'):
        kl_count = TERMWEEKS // 2
        kl_start = TERMWEEKS // 2
    elif(klass['ClassWeeks'] == '单'):
        kl_count = TERMWEEKS // 2
        kl_interval = 2
    elif(klass['ClassWeeks'] == '双'):
        kl_count = TERMWEEKS //2
        kl_interval = 2
        kl_start = 1
    else:
        (kl_start, kl_end) = klass['ClassWeeks'].split('-')
        kl_start = int(kl_start) - 1
        kl_count = int(kl_end) - kl_start

    event.add('uid', str(uuid1())+'@THU')
    event.add('summary', klass['Summary'])
    event.add('dtstart', datetime(*klass['TimeStart']) + timedelta(kl_start * 7))
    event.add('dtend',   datetime(*klass['TimeEnd']) + timedelta(kl_start * 7))
    event.add('dtstamp', datetime.today())
    event.add('location',klass['Location'])
    event.add('rrule', {'freq': 'weekly', 'interval': kl_interval, 'count': kl_count})
    calendar.add_component(event)

with open('classtable.ics', 'w+', encoding='utf-8') as ics:
    ics.write(calendar.to_ical().decode('utf-8'.replace('\r\n', '\n').strip()))


import numpy as np
import matplotlib.pyplot as plt
import pandas
'''
R is the number of activated uninfected cells
L is the number of latently infected cells
E is the number of activated infected cells
V is the number of free virions
gamma is the rate at which uninfected cells arise
tau is the proportion of cells activated
mu is the HIV-independent death rate of uninfected cells
beta is the rate of infection of cells per virion
rho is the proportion of latently infected cells upon infection
alpha is the activation rate of latently infected cells
delta is the death rate of actively infected cells
pi is the rate of porduction of virions by an actively infected cells
sigma is the removal rate of cell-free virus
'''

sim_length = 120 # days
# parameters that can change - initial values given by paper
total_cells = 1000
gamma = 1.36
mu = 1.36 * 10**-3
tau = 0.2
beta = 0.00027
rho = 0.1
alpha = 3.6 * 10**-2
sigma = 2
delta = 0.33
pi = 100
deltat = 0.75 # days

class Plasma:
    def __init__(self):
        self.R = total_cells * tau
        self.L = 0
        self.E = 0
        self.V = 4 * 10**-7
        #self.unactivated_cells = total_cells - self.R - self.L - self.E
        print('initialized')

    def update(self, deltat):
        print('updating')
        global total_cells
        dRdt = gamma * tau - mu * self.R - beta * self.R * self.V
        dLdt = rho * beta * self.R * self.V - mu * self.L - alpha * self.L
        dEdt = (1 - rho) * beta * self.R * self.V + alpha * self.L - delta * self.E
        dVdt = pi * self.E - sigma * self.V
        print('dRdt {} dLdt {} dEdt {} dVdt {}'.format(dRdt, dLdt, dEdt, dVdt))

        self.R += dRdt * deltat
        self.L += dLdt * deltat
        self.E += dEdt * deltat
        self.V += dVdt * deltat
        print("values: R {}, L {}, E {}, V {}".format(self.R, self.L, self.E, self.V))
        return self.R, self.L, self.E, self.V

def main():
    num_steps = sim_length / deltat
    print(num_steps)
    plasma = Plasma()
    graph_array = [[-1, plasma.R, plasma.L, plasma.E, plasma.V]]
    time_list = []
    R_list = []
    L_list = []
    E_list = []
    V_list = []
    for i in range(int(num_steps)):
        R, L, E, V = plasma.update(deltat)
        time_list.append(i)
        R_list.append(R)
        L_list.append(L)
        E_list.append(E)
        V_list.append(V)
        graph_array.append([i, R, L, E, V])
    for x in time_list:
        time_list[x] = time_list[x] * deltat
    array_collection = [time_list, R_list, L_list, E_list, V_list]
    # code below adapted from : https://matplotlib.org/gallery/api/two_scales.html
    fig, ax1 = plt.subplots()
    color = 'tab:red'
    ax1.set_xlabel('Time (days)')
    ax1.set_ylabel('R', color=color)
    ax1.plot(array_collection[0], array_collection[1], color=color)
    ax1.tick_params(axis='y', labelcolor=color)
    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis
    color = 'tab:blue'
    ax2.set_ylabel('L and E', color=color)  # we already handled the x-label with ax1
    ax2.semilogy(array_collection[0], array_collection[2], color='blue', label='L')
    ax2.semilogy(array_collection[0], array_collection[3], color='green', label='E')
    ax2.tick_params(axis='y', labelcolor=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    plt.legend()
    plt.show()

if __name__ == '__main__':
    main()


from setuptools import setup, find_packages
from os import path

# read the contents of your README file
this_directory = path.abspath(path.dirname(__file__))
with open(path.join(this_directory, 'README.md'), encoding='utf-8') as f:
    long_description = f.read()

setup(
    name='ctapipe_io_targetio',
    packages=find_packages(),
    version='0.1',
    description='ctapipe plugin for reading target files',
    long_description=long_description,
    long_description_content_type='text/markdown',
    install_requires=[
        'astropy',
        'ctapipe',
    ],
    tests_require=['pytest'],
    setup_requires=['pytest_runner'],
    author='Jason Watson',
    author_email='jason.jw@live.co.uk',
    license='MIT',
)