# -*- coding: utf-8 -*-
# 所有Handler基类

import copy

import tornado.web

class BaseHandler(tornado.web.RequestHandler):
    """所有Handler基类.

    重载了输出方法，
    正确输出 (200) 使用 write_json_f 方法，错误输出使用 write_error_f 方法
    """

    # 输出模板
    RSP = {
        "code": "200",
        "message": "OK",
        "relationships": {
            "author": "null"
        },
        "jsonapi": {
            "version": "null"
        }}

    # 输出JSON
    rsp = dict

    pagination = None

    _ARG_DEFAULT = object()

    def get_argument_t(self, name, default=_ARG_DEFAULT, type=None, strip=True):
        """重载get_argument方法，并判断参数类型是否匹配.

        输出时如果满足参数类型要求，则按类型输出，否则输出None
        """
        arg = self.get_argument(name, default, strip)
        if type == int:
            if not isinstance(arg, int):
                arg = int(arg) if arg.isdigit() else None
        return arg

    def get_argument_cardno(self, default=0):
        """获取cardno参数，并校验."""
        cardno = self.get_argument_t('cardno', default, int)
        if not cardno or cardno < 213120000 or cardno > 213259999:
            return self.write_error(4001)
        return cardno

    def write_json_f(self, data=None):
        """重载write方法，并finish.

        Args:
            data: 输出JSON的data字段，默认为空
        """
        self.write_json(data)
        self.finish()

    def write_json(self, data=None):
        """重载write方法.

        Args:
            data: 输出JSON的data字段，默认为空
        """
        self.rsp = copy.deepcopy(self.RSP)
        if self.pagination:
            self.rsp["pagination"] = self.pagination
        self._write_json(data)

    def _write_json(self, data=None):
        """输出准备."""
        if data:
            self.rsp["data"] = data
        self.rsp["relationships"]["author"] = self.INFO["author"]
        self.rsp["jsonapi"]["version"] = self.INFO["version"]
        self.write(self.rsp)

    def write_error_f(self, status_code):
        """重载write_error方法，并finish.

        Args:
            status_code: HTTP状态码
        """
        self.write_error(status_code)
        self.finish()

    def write_error(self, status_code, **kwargs):
        """重载write_error方法.

        Args:
            status_code: HTTP状态码
        """
        self.rsp = copy.deepcopy(self.RSP)
        self._transforms = []
        self._write_error(status_code)

    def _write_error(self, status_code):
        """write_error输出的准备.

        Args:
            status_code: HTTP状态码，根据此参数配置输出内容
        """
        status_code = str(status_code)
        status_code_b, status_code_s = status_code[:3], status_code[3:]

        self.rsp["code"] = status_code_b

        response = {
            "400": {
                "message": "Bad Request",
                "errors": {
                    '1': {"code": 4001,
                          "message": "Wrong Type Of Cardno"}
                }
            },
            "405": {
                "message": "Method Not Allowed"
            },
            "500": {
                "message": "Internal Server Error",
                "errors": {
                    '1': {"code": 5001,
                          "message": "Unknown Server Error"}
                }
            }
        }
        self.set_status(int(status_code_b))
        self.rsp["message"] = response.get(status_code_b, {}).get(
            "message", "Undefined Error")
        self.rsp["errors"] = response.get(status_code_b, {}).get(
            "errors", {}).get(status_code_s, {"code": status_code,
                                              "message": "Undefined Error"})

        self._write_json()


#!/usr/bin/python

# Test main target alternatives.

from BoostBuild import Tester
from string import find
t = Tester()

# Test that basic alternatives selection works.
t.write("project-root.jam", " ")
t.write("Jamfile", """

exe a : a_empty.cpp ;
exe a : a.cpp : <variant>release ;
""")
t.write("a_empty.cpp", "")
t.write("a.cpp", "int main() { return 0; }\n")

t.run_build_system("release")
t.expect_addition("bin/$toolset/release/a.exe")

# Test that alternative selection works for ordinary
# properties, in particular user-defined.
t.write("project-root.jam", " ")
t.write("Jamfile", """

import feature ;
feature.feature X : off on : propagated ;

exe a : b.cpp ;
exe a : a.cpp : <X>on ;
""")
t.write("b.cpp", "int main() { return 0; }\n")

t.rm("bin")

t.run_build_system()
t.expect_addition("bin/$toolset/debug/b.obj")

t.run_build_system("X=on")
t.expect_addition("bin/$toolset/debug/X-on/a.obj")

t.rm("bin")

# Test that everything works ok even with default
# build.

t.write("Jamfile", """

exe a : a_empty.cpp : <variant>release ;
exe a : a.cpp : <variant>debug ;
""")

t.run_build_system()
t.expect_addition("bin/$toolset/debug/a.exe")

# Test that only properties which are in build request
# matters when selection alternative. IOW, alternative
# with <variant>release is better than one with
# <variant>debug when building release version.
t.write("Jamfile", """

exe a : a_empty.cpp : <variant>debug ;
exe a : a.cpp : <variant>release ;
""")

t.run_build_system("release")
t.expect_addition("bin/$toolset/release/a.exe")

# Test that free properties do not matter. We really don't
# want <cxxflags> property in build request to affect
# alternative selection.
t.write("Jamfile", """
exe a : a_empty.cpp : <variant>debug <define>FOO <include>BAR ;
exe a : a.cpp : <variant>release ;
""")

t.rm("bin/$toolset/release/a.exe")
t.run_build_system("release define=FOO")
t.expect_addition("bin/$toolset/release/a.exe")

# Test that abibuity is reported correctly
t.write("Jamfile", """
exe a : a_empty.cpp ;
exe a : a.cpp ;
""")
t.run_build_system("--no-error-backtrace", status=1)
t.fail_test(find(t.stdout(), "because no best-matching alternative could be found") == -1)

# Another ambiguity test: two matches properties in one alternative are
# neither better nor worse than a single one in another alternative.
t.write("Jamfile", """
exe a : a_empty.cpp : <optimization>off <profiling>off ;
exe a : a.cpp : <debug-symbols>on ;
""")

t.run_build_system("--no-error-backtrace", status=1)
t.fail_test(find(t.stdout(), "because no best-matching alternative could be found") == -1)

# Test that we can have alternative without sources
t.write("Jamfile", """
alias specific-sources ;
import feature ;
feature.extend os : MAGIC ;
alias specific-sources : b.cpp : <os>MAGIC ;
exe a : a.cpp specific-sources ;
""")
t.rm("bin")
t.run_build_system()

t.cleanup()


__author__ = "Andrew DiFabbio"
__email__ = "avd38@cornell.edu"

import argparse
import numpy as np
import pandas as pd
import os
from matplotlib import pyplot as plt
import exp_params as exp
from scipy.stats import chi2

np.seterr(divide='ignore', invalid='ignore')

# COMBINING REDUCED DATA----------------------------------------------------------------------------------------
def sum_over_results(path,num_ch,num_bins,verbose=False):
    """Sum the coincidences, accidentals, and fluorescence counts for a thickness within a
    folder, returning the result in the format <fluorescence counts>,<coincidences>,
    <accidentals>.

    If verbose, these results are returned with respect to every permuation of fluorescence
    and scattered detector. Otherwise they are simply returned according to fluorescence
    detector."""

    fl_ct=np.zeros((num_ch,1))
    coin=np.zeros((num_ch*num_ch,num_bins))
    acoin=np.zeros((num_ch*num_ch,num_bins))

    for l in os.listdir(path):
        if '.txt' in l: continue

        if 'fluor_count' in l:
            fl_ct+=np.genfromtxt(os.path.join(path,l),delimiter=',').reshape((4,1))
        elif 'coincidence' in l:
            coin+=np.genfromtxt(os.path.join(path,l),delimiter=',')
        elif 'accidentals' in l:
            acoin+=np.genfromtxt(os.path.join(path,l),delimiter=',')

    if not verbose:
        coin=np.sum(coin.reshape((num_ch,num_ch,num_bins)),axis=1)
        acoin=np.sum(acoin.reshape((num_ch,num_ch,num_bins)),axis=1)

    return fl_ct,coin,acoin

def get_ang_corr(num_ch):
    """Calculate the factor for angular dependence in IAB as it relates to the detectors. This
    is essentially an integral over the solid angle of the angular coefficient."""

    ang_corr=np.zeros((num_ch,1))
    for r in range(num_ch):
        for ch in range(num_ch):
            if r==ch: continue
            ang_corr[r]+=exp.CHANNEL_ANGULAR_COEFF[ch]*exp.CHANNEL_SOLID_ANGLE[ch]

    return np.power(ang_corr,-1)

def combine_summed_results(data,num_ch,num_orb,e_bin,n_bins,atten=True,thick=None):
    """Combine the already summed counts by applying angular and attenuation corrections and
    handling error propagation as necessary. Returns <normalized excess>,<normalized error>.

    If atten, attenuation correcion is applied, and a thickness must be given in the form
    'data/40nm/', for example."""

    if atten and thick == None: raise Exception('combine_summed_results: attenuation correction needs thickness.')
    ang_corr=get_ang_corr(num_ch)
    fl_ct,coin,acoin=data

    excesses=ang_corr*(coin-acoin/num_orb)/fl_ct
    variances=ang_corr**2*((coin+acoin/num_orb**2)/fl_ct**2 + (coin-acoin/num_orb)**2/fl_ct**3)
    if atten: excesses,variances=correct_atten(excesses,variances,e_bin,n_bins,thick)

    norm_vars=np.power(np.sum(np.power(variances,-1),axis=0),-1)

    return (norm_vars*np.sum(excesses/variances,axis=0))/(e_bin*exp.XMAP_ENERGY_UNIT), \
        np.sqrt(norm_vars)/(e_bin*exp.XMAP_ENERGY_UNIT)

def combine_folder(path,num_ch,num_bins,num_orb,e_bin,atten=True,thick=None):
    """Read and combine the data for a given thickness."""

    if atten and thick == None: raise Exception('combine_folder: attenuation correction needs thickness.')
    return combine_summed_results(sum_over_results(path,num_ch,num_bins),num_ch,num_orb,e_bin,num_bins,atten,thick)

def combine_data(wdir,num_ch,num_bins,num_orb,e_bin,atten=True,write=True):
    """Fully combine the reduced data. If atten, attenuation correction is applied. If write, data
    will be written to CSV files afterward, otherwise the result is returned. The results consist
    of normalized excess and corresponding errors with rows representing thickness in the order 40nm,
    80nm, 160nm, 320nm, empty, and columns representing energy bin."""

    result_excesses=np.zeros((5,num_bins))
    result_errors=np.zeros((5,num_bins))

    for r,t in enumerate(['data/40nm/','data/80nm/','data/160nm/','data/320nm/','data/empty/']):
        exc,err=combine_folder(os.path.join(wdir,t),num_ch,num_bins,num_orb,e_bin,atten,t)
        result_excesses[r,:]+=exc
        result_errors[r:]+=err

    if write:
        np.savetxt(os.path.join(wdir,'combined_result_excesses.csv'),result_excesses,delimiter=',')
        np.savetxt(os.path.join(wdir,'combined_result_errors.csv'),result_errors,delimiter=',')
    else:
        return result_excesses,result_errors

# PLOTTING COMBINED RESULTS AND PREDICTIONS---------------------------------------------------------------------
def k(t):
    return exp.K_COEFF*(t/exp.D)**0.75

def W_secondaries(t,Es):
    return 2*k(t)*exp.Z*(exp.E0-exp.EB-Es)/(4*np.pi*Es)

def W_LET(Es):
    return 2*exp.ALPHA*((exp.E0-exp.EB)/exp.MC2)/(4*np.pi**2*Es)

def plot_thickness(ax,exc,err,Emin,Emax,e_bin,t):
    """Plot the combined results for a given thickness and energy range along with the predicted rates
    for that thickness. Blue for secondaries, red for LET, and green for secondaries+LET."""

    start,stop=int(Emin/(e_bin*exp.XMAP_ENERGY_UNIT)),int(Emax/(e_bin*exp.XMAP_ENERGY_UNIT)+0.5)
    Erange=np.arange((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                     (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),e_bin*exp.XMAP_ENERGY_UNIT)
    Erange_=np.linspace((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                      (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),1000)
    # plot predictions
    ax.plot(Erange_,0*Erange_,color='k')
    ax.plot(Erange_,W_secondaries(t,Erange_),color='b')
    ax.plot(Erange_,W_LET(Erange_),color='r')
    ax.plot(Erange_,W_LET(Erange_)+W_secondaries(t,Erange_),color='g')

    # plot excesses
    ax.errorbar(Erange,exc[start:stop],yerr=err[start:stop],ls='none')
    ax.scatter(Erange,exc[start:stop])
    ax.set_title(f'{t}nm')
    ax.set_xlabel('Energy (keV)')
    ax.set_ylabel('Normalized Excess Rate')

def plot(wdir,Emin,Emax,e_bin,save=False):
    """Plot the combined results and predictions for 40nm,80nm,160nm,320nm thicknesses in a given energy
    range."""

    excess=np.genfromtxt(os.path.join(wdir,'combined_result_excesses.csv'),delimiter=',')
    errors=np.genfromtxt(os.path.join(wdir,'combined_result_errors.csv'),delimiter=',')

    fig,((ax0,ax1),(ax2,ax3))=plt.subplots(2,2)
    plot_thickness(ax0,excess[0,:],errors[0,:],Emin,Emax,e_bin,40)
    plot_thickness(ax1,excess[1,:],errors[1,:],Emin,Emax,e_bin,80)
    plot_thickness(ax2,excess[2,:],errors[2,:],Emin,Emax,e_bin,160)
    plot_thickness(ax3,excess[3,:],errors[3,:],Emin,Emax,e_bin,320)
    plt.tight_layout()

    if save:
        if not os.path.isdir(os.path.join(wdir,'figures/')): os.mkdir(os.path.join(wdir,'figures/'))
        plt.savefig(os.path.join(os.path.join(wdir,'figures/'), \
                                 f'normalized_excess_plots_(Emin={Emin},Emax={Emax}).png'))
    else:
        plt.show()

# CALCULATE AND REPORT CHI SQUARE VALUES------------------------------------------------------------------------
def calc_chisq(obs,err,pred,rnd=True):
    """Calculate the chi square and probability given observed results, errors, and predicted results.
    If rnd, round the returned values accordingly."""

    chi_sq=np.sum((obs-pred)**2/err**2)
    nu=obs.shape
    p=(1-chi2.cdf(chi_sq,nu))[0]
    if rnd: return round(chi_sq,3),round(p,5)
    return chi_sq,p

def excess_chisq(wdir,Emin,Emax,e_bin,save=False):
    """Compute the chi square and probability between observed and predicted results of the combined data
    given an energy range. Print the results as a table to the command line and write them to a TXT"""

    excess=np.genfromtxt(os.path.join(wdir,'combined_result_excesses.csv'),delimiter=',')
    errors=np.genfromtxt(os.path.join(wdir,'combined_result_errors.csv'),delimiter=',')
    start,stop=int(Emin/(e_bin*exp.XMAP_ENERGY_UNIT)),int(Emax/(e_bin*exp.XMAP_ENERGY_UNIT)+0.5)
    Erange=np.arange((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                     (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),e_bin*exp.XMAP_ENERGY_UNIT)

    table=[('','df','SEC chi2','SEC p','LET chi2','LET p','TOT chi2','TOT p')]
    df=(Erange.shape)[0]
    for r,t in enumerate([40,80,160,320]):
        SECchi2,SECp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_secondaries(t,Erange))
        LETchi2,LETp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_LET(Erange))
        TOTchi2,TOTp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_LET(Erange)+W_secondaries(t,Erange))
        table.append((str(s) for s in (f'{t}nm',df,SECchi2,SECp,LETchi2,LETp,TOTchi2,TOTp)))

    if save:
        if not os.path.isdir(os.path.join(wdir,'chisq/')): os.mkdir(os.path.join(wdir,'chisq/'))
        with open(os.path.join(os.path.join(wdir,'chisq/'),f'chisq_table_(Emin={Emin},Emax={Emax}).txt'),'w+') as f:
            for row in table: f.write('| {:>6} | {:>2} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} |\n'.format(*row))
    for row in table: print('| {:>6} | {:>2} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} |'.format(*row))

# CALCULATE AND PRESENT CORRECTIONS FOR ATTENUATION-------------------------------------------------------------
def get_atten(Erange,n_ch=4,verbose=False):
    """Compute the attentuation from all sources for a given energy range. If verbose, the result includes the
    raw transmittance for every source of attenuation in the format <He>,<kapton>,<1-elem detector QE>,<4-elem
    detector QE>,<320nm>,<160nm>,<80nm>,<40nm>.
    Otherwise, it returns transmittance results for each mode (based on the detector channel that saw the
    fluorescence event) and for each target thickness relative to the transmittance at fluorescence energy.
    These are returned in the format <average mode transmittance>,(<320nm>,<160nm>,<80nm>,<40nm>).

    This function, along with other functions related to attenuation correction, will likely need to be modified
    if any of the experimental conditions change. This is far less generalizable than everything else."""

    atten_raw=pd.read_csv(exp.ATTENUATION_DOC_PATH)
    He_T=np.interp(Erange,atten_raw['He_E'].values,atten_raw['He_T'].values)
    Kap_T=np.interp(Erange,atten_raw['Kap_E'].values,atten_raw['Kap_T'].values)
    D1_T=np.interp(Erange,atten_raw['D_E'].values,atten_raw['D1_T'].values)
    D4_T=np.interp(Erange,atten_raw['D_E'].values,atten_raw['D4_T'].values)
    M320_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M320_T'].values)
    M160_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M160_T'].values)
    M80_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M80_T'].values)
    M40_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M40_T'].values)

    if verbose: return He_T,Kap_T,D1_T,D4_T,M320_T,M160_T,M80_T,M40_T

    D1_Tot=He_T*Kap_T*D1_T
    D4_Tot=He_T*Kap_T*D4_T
    Mode_T=np.zeros((n_ch,len(Erange)))

    for r in range(n_ch):
        S=0
        for ch in range(n_ch):
            if r==ch: continue
            S+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]
            if ch == 0:
                Mode_T[r,:]+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]*D1_Tot
            else:
                Mode_T[r,:]+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]*D4_Tot

        Mode_T[r,:]/=S
        if r == 0:
            Mode_T[r,:]/=np.interp(exp.E_FLUOR,Erange,D1_Tot)
        else:
            Mode_T[r,:]/=np.interp(exp.E_FLUOR,Erange,D4_Tot)

    M320_T/=np.interp(exp.E_FLUOR,Erange,M320_T)
    M160_T/=np.interp(exp.E_FLUOR,Erange,M160_T)
    M80_T/=np.interp(exp.E_FLUOR,Erange,M80_T)
    M40_T/=np.interp(exp.E_FLUOR,Erange,M40_T)

    return Mode_T,(M320_T,M160_T,M80_T,M40_T)

def correct_atten(exc,err,e_bin,n_bins,t):
    """Similarly to how corrections are applied for the angular component, this function corrects the excesses
    and energies for attenuation."""

    Mode_T,M_T=get_atten(exp.XMAP_ENERGY_UNIT*e_bin*np.arange(n_bins),4)
    th={'data/40nm/':3,'data/80nm/':2,'data/160nm/':1,'data/320nm/':0}

    if t == 'data/empty/': return exc/Mode_T,err/Mode_T
    return exc/(Mode_T*M_T[th[t]]),err/(Mode_T*M_T[th[t]])

def plot_atten(Emin,Emax,n_ch=4,save=False,wdir=None):
    """Plot the transmittance for all sources of attenuation over a given energy range. If save, save the result
    as a figure."""

    if save and wdir == None: raise Exception('Instructed to save with no wdir provided.')

    Erange=np.linspace(Emin,Emax,1000)
    He_T,Kap_T,D1_T,D4_T,M320_T,M160_T,M80_T,M40_T=get_atten(Erange,n_ch,True)

    plt.plot(Erange,M40_T,label='Material Transmittance 40nm')
    plt.plot(Erange,M80_T,label='Material Transmittance 80nm')
    plt.plot(Erange,M160_T,label='Material Transmittance 160nm')
    plt.plot(Erange,M320_T,label='Material Transmittance 320nm')
    plt.plot(Erange,He_T,label='Helium Transmittance')
    plt.plot(Erange,Kap_T,label='Kapton Transmittance')
    plt.plot(Erange,D1_T,label='1-Elem Detector QE')
    plt.plot(Erange,D4_T,label='4-Eleme Detector QE')

    plt.legend(loc='lower right')
    plt.title('Sources of Attenuation')
    plt.xlabel('Energy (keV)')
    plt.ylabel('Transmittance')
    plt.tight_layout()

    if save:
        if not os.path.isdir(os.path.join(wdir,'figures/')): os.mkdir(os.path.join(wdir,'figures/'))
        plt.savefig(os.path.join(os.path.join(wdir,'figures/'), \
                                 f'atten_sources_(Emin={Emin},Emax={Emax}).png'))
    else:
        plt.show()

# HANDLING COMMAND LINE ARGUMENTS-------------------------------------------------------------------------------
if __name__=='__main__':
    # Parsing command line arguments
    p = argparse.ArgumentParser()
    sp=p.add_subparsers(dest='command')

    # combine command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py combine <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <NUMBER OF BINS> <NUMBER OF CHANNELS> <NUMBER OF ORBITS> <--skipwrite> <--noatten>"""

    combine_p=sp.add_parser('combine',help='combine reduced data')
    combine_p.add_argument(help='working directory',dest='wdir')
    combine_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    combine_p.add_argument(help='number of bins in final histogram',type=int,dest='n_bins')
    combine_p.add_argument(help='number of detector channels',type=int,dest='n_ch')
    combine_p.add_argument(help='number of orbits over which to average accidentals',type=int,dest='n_orb')
    combine_p.add_argument('--skipwrite',help='do not write the result to a CSV',action='store_false')
    combine_p.add_argument('--noatten',help='do not correct for attenuation',action='store_false')

    # plot command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py plot <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> <--save>"""

    plot_p=sp.add_parser('plot',help='plot excesses command')
    plot_p.add_argument(help='working directory',dest='wdir')
    plot_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    plot_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    plot_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    plot_p.add_argument('--save',help='save the plot generated',action='store_true')

    # chisq command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py chisq <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> <--save>"""

    chisq_p=sp.add_parser('chisq',help='calculate chi square results')
    chisq_p.add_argument(help='working directory',dest='wdir')
    chisq_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    chisq_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    chisq_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    chisq_p.add_argument('--save',help='save the plot generated',action='store_true')

    # pltatten command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py pltatten <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> \\
    <--wdir PATH TO WORKING DIRECTORY> <--save>"""

    pltatten_p=sp.add_parser('pltatten',help='plot sources of attenuation')
    pltatten_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    pltatten_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    pltatten_p.add_argument('--wdir',help='working directory for saving the plot',default=None)
    pltatten_p.add_argument('--save',help='save the plot generated',action='store_true')

    a=p.parse_args()
    if a.command=='plot':
        plot(a.wdir,a.Emin,a.Emax,a.e_bin,a.save)
    elif a.command=='combine':
        combine_data(a.wdir,a.n_ch,a.n_bins,a.n_orb,a.e_bin,a.noatten,a.skipwrite)
    elif a.command=='chisq':
        excess_chisq(a.wdir,a.Emin,a.Emax,a.e_bin,a.save)
    elif a.command=='pltatten':
        plot_atten(a.Emin,a.Emax,4,a.save,a.wdir)


from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec

from pyshell.util import ipydb
from pyshell.nosetests.postmortem import PostMortemScript
ipydb()

import glob
import os.path
import re

class SlopeDisplay(PostMortemScript):
    """Do the slope display."""

    module_ids = [ 'tests.test_zernike.test_zernike_slopes' ]

    filename_re = re.compile(r'([\w]+)_([\w]+)_([\d]+)_([\d-]+)')

    def parse_filename(self, filename):
        """docstring for parse_filename"""
        match = self.filename_re.match(filename)
        groups = match.groups()

        info = {
            'axes' : groups[0],
            'type' : groups[1],
            'n' : int(groups[2]),
            'm' : int(groups[3])
        }

        return info

    def do(self):
        """Actually do the hard work."""
        for key in self.test_data[self.module_ids[0]].keys():
            self.do_test(key)
        plt.show()

    def do_test(self, test):
        """Make displays for a specific test case."""
        filedata = {}
        filenames = self.test_data[self.module_ids[0]][test].keys()
        for filename in filenames:
            filedata[filename] = self.parse_filename(filename)

        row = { 'xs': 0, 'ys': 1 }
        col = { 'npy' : 1, 'pyt': 2, 'dif':3}
        gs = gridspec.GridSpec(2,5)
        figure = plt.figure(figsize=(8,4))

        for filename in filenames:
            info = filedata[filename]
            data = self.test_data[self.module_ids[0]][test][filename]
            col_pos = col.get(info['type'],0)
            row_pos = row.get(info['axes'],0)
            if info['axes'] == 'z':
                col_pos = 0
            axes = self.axes(figure, gs[row_pos,col_pos])
            self.label_axes(axes, "{:s}-{:s}-({:2d},{:2d})".format(
                info['axes'][0].capitalize(), info['type'], info['n'], info['m']
            ))
            axes.imshow(data)

    def axes(self, figure, subplotspec):
        """Return a setup axes from a subplotspec"""
        ax = figure.add_subplot(subplotspec)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        ax.set_frame_on(False)
        return ax

    def label_axes(self, ax, label):
        """Label axes above."""
        ax.text(0.5, 1.0, label, transform=ax.transAxes, va='bottom', ha='center')

if __name__ == '__main__':
    SlopeDisplay.script()


# task_1

month = 1

for remain in range(100, -1, -4):
    print(f'В {month} месяце у вас будет {remain} килограмм гречки')
    month += 1

print('Запасы закончились')

# task_2

count = int(input('Введите количество должников: '))
sum_debt = 0

for debtor in range(0, count, 5):
    print(f'Должник с номером {debtor}')
    debt = int(input('Сколько должны? '))
    sum_debt += debt

print(f'Общая сумма долга: {sum_debt}')

# task_3

timer = int(input('Введите количество секунд: '))

for reverse_time in range(timer, 0, -1):
    end = int(input('Закончить? '))
    if end == 1:
        print(f'Ваша еда готова, можете забрать. Оставшееся время было: {reverse_time}')
        break
else:
    print('Ваша еда готова, осторожно горячo!')

# task_4

left = int(input('Введите левую границу: '))
right = int(input('Введите правую границу: '))
divider = int(input('Введите число, на которое должно делиться другое число без остатка: '))

num_sum = 0
num_count = 0

for num in range(left, right + 1):
    if num % divider == 0:
        num_sum += num
        num_count += 1

print(f'Среднее арифметическое: {num_sum / num_count}')

# task_5

start = int(input('Введите начало отрезка: '))
end = int(input('Введите конец отрезка: '))
step = int(input('Введите шаг (отрицательный): '))

for num in range(end, start - 1, step):
    print(f'В точке {num} функция равна {num ** 3 + 2 * num ** 2 - 4 * num + 1}')

# task_6

educational_grant = int(input('Введите стипендию: '))
expenses = int(input('Введите расходы на проживание: '))

for month in range(1, 11):
    print(f'{month} месяц. Траты {expenses}, не хватает {expenses - educational_grant}')

    educational_grant *= 1.03
    expenses *= 1.03

print(f'Нужно попросить у родителей {expenses - educational_grant} рублей')

# task_7

n = int(input('Введите N: '))
res = 0

for i in range(n):
    res += -1 ** i * 1 / 2 ** i

print(res)

# task_8

boys = int(input('Введите количество мальчиков: '))
girls = int(input('Введите количество девочек: '))
people = boys + girls

if boys == 0 and girls == 0:
    print('')
elif boys == 1 and girls == 0:
    print('B')
elif boys == 0 and girls == 1:
    print('G')
elif boys == 1 and girls == 1:
    print('BG')
elif boys / 2 > girls or girls / 2 > boys:
    print('Решения нет')
else:
    res = ''
    last = ''

    if girls > boys:
        res = 'GB'
        last = 'B'
    else:
        res = 'BG'
        last = 'G'
    count = 2
    girls -= 1
    boys -= 1

    while count < people:
        if girls > boys and last != 'GG':
            if res[-1] == 'B':
                last = ''
            res += 'G'
            last += 'G'
            girls -= 1
        elif boys > girls and last != 'BB':
            if res[-1] == 'G':
                last = ''
            res += 'B'
            last += 'B'
            boys -= 1
        else:
            if last == 'GG' or last == 'G':
                res += 'B'
                last = 'B'
                boys -= 1
            else:
                res += 'G'
                last = 'G'
                girls -= 1
        count += 1
    print(res)


from django.urls import path
from . import views

app_name = 'quiz_urls'

urlpatterns = [
    path('', views.index_quiz, name='index_quiz'),
    path('<quiz_pk>/hard', views.quiz_questions, name='quiz_detail_hand'),
    path('<quiz_pk>/easy', views.quiz_questions_paginate, name='quiz_detail_easy'),
    path('<quiz_pk>/quiz_finish', views.quiz_finish, name='quiz_finish'),
    path('<quiz_pk>/restart_test', views.restart_test, name='restart_test'),
]


from torch.utils.data import DataLoader
from datasets import SEMDataset
from model_elmo import Model
import argparse
import torch
import torch.nn as nn
from tqdm import tqdm
import torch.utils.data as data
from evaluate import eval
import os
import warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def collate_fn(batch):
    # put question lengths in descending order so that we can use packed sequences later
    batch.sort(key=lambda x: x[-1], reverse=True)
    return data.dataloader.default_collate(batch)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=80)
    parser.add_argument('--num_hid', type=int, default=2048)#512,2048(elmo)
    parser.add_argument('--emb_dim', type=int, default=2048)#300,2048(elmo)
    parser.add_argument('--tag_emb_dim', type=int, default=100)
    parser.add_argument('--pos_emb_dim', type=int, default=512)#300,512(elmo)
    parser.add_argument('--wordG_out_dim', type=int, default=1024)  # 512,1024(transformer)
    parser.add_argument('--output', type=str, default='saved_models/')
    parser.add_argument('--idx2word', type=str, default='data/idx2word.pkl')
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--ntoken', type=int, default=5342)
    parser.add_argument('--ntag', type=int, default=46)
    parser.add_argument('--npos', type=int, default=74)
    parser.add_argument('--num_class', type=int, default=3)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--options_file', type=str, default="https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/elmo_2x4096_512_2048cnn_2xhighway_options.json")
    parser.add_argument('--weight_file', type=str, default="https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/elmo_2x4096_512_2048cnn_2xhighway_weights.hdf5")
    parser.add_argument('--alpha', type=float, default=0.2, help='Alpha for the leaky_relu.')
    parser.add_argument('--seed', type=int, default=1111, help='random seed')
    parser.add_argument('--model_name', type=str, default="model_elmo", help='saved models name')
    args = parser.parse_args()
    return args

def train(model,train_loader,dev_loader,args):
    optim = torch.optim.Adam(model.parameters(),lr=0.0001)
    log_softmax = nn.LogSoftmax(dim=2).to(device)
    div_loss = nn.KLDivLoss(reduction='none')
    softmax = nn.Softmax(dim=2).to(device)
    best_socre = 0
    best_epoch = -1
    for epoch in range(args.epochs):
        train_tq = tqdm(train_loader,desc='E{:03d}'.format(epoch),ncols=0)
        dev_tq = tqdm(dev_loader,desc='T{:03d}'.format(epoch),ncols=0)
        total_loss = 0
        model.eval()
        matchm,rankm ,score , acc, o_acc, b_acc,i_acc = evaluate(model, dev_tq,softmax,epoch)
        print(rankm,acc, o_acc, b_acc,i_acc)
        if score>best_socre:
            best_socre = score
            best_epoch = epoch
            best_match = matchm
            torch.save(model.state_dict(),args.model_name+'.pkl')

        model.train()
        for (s,a,p,a_s,s_len,mask,mask2,pos_len,graph,indicate,posid,distribution) in train_tq:
            s = s.to(device) #[B, 38]
            a = a.to(device) #[B, 38]
            p = p.to(device) #[B,38]
            a_s = a_s.to(device)#[B,9,38]
            s_len = s_len.to(device)#[B]
            mask = mask.to(device)
            mask2 = mask2.to(device)
            pos_len = pos_len.to(device)
            graph = graph.to(device)
            #indicate = indicate.to(device)
            posid = posid.to(device)
            distribution = distribution.to(device)

            prob = model(s,p,mask2,posid,graph,indicate) #[B,38,3]
            prob = log_softmax(prob)

            # div loss
            #distribution = distribution/(torch.sum(distribution,dim=2,keepdim=True)+0.00001)
            #div_l = div_loss(prob,distribution)
            _, p_max = prob.max(dim=2)

            #total = s_len.sum()
            #loss = div_l.sum()/total

            weight = torch.ones(prob.shape[0],38,3).to(device)
            weight[:,:,1:3]=5
            clamp_a = torch.clamp(a,0,3)
            #-----------------------
            #calculate loss
            new_a = torch.unsqueeze(clamp_a,-1)
            new_b = torch.zeros(prob.shape[0],38,args.num_class).to(device)
            new_b.scatter_(2, new_a, 1)
            total = s_len.sum()
            loss = (new_b*mask*prob).sum()
            loss = -loss/total
            #------------------------

            total_loss += loss.item()*total
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 2)
            optim.step()
            optim.zero_grad()
            fmt = '{:.4f}'.format
            train_tq.set_postfix(loss=fmt(loss.item()))
        print("total loss:",total_loss)

    print(best_socre)
    print(best_match)
    print(best_epoch)

def evaluate(model,dev_tq,softmax,epoch):

    total_o = 0
    total_b = 0
    total_i = 0
    o_acc = 0
    b_acc = 0
    i_acc = 0
    acc = 0
    total = 0
    with torch.no_grad():
        for (s,a,p,a_s,s_len,mask,mask2,pos_len,graph,indicate,posid,distribution) in dev_tq:

            s = s.to(device)  # [B, 38]
            p = p.to(device)
            #mask = mask.to(device)
            mask2 = mask2.to(device)

            pos_len = pos_len.to(device) #[64]
            graph = graph.to(device)    #[64,80,80]
            #indicate = indicate.to(device)#[64,80]
            posid = posid.to(device) #[64,80]

            prob = model(s,p,mask2,posid,graph,indicate) #[dev_set,38,3]
            prob = softmax(prob)
            prob = prob.cpu()
            ac, o, b, i, num_o, num_b, num_i, num = accuracy(prob,distribution,mask[:,:,0])
            acc+=ac
            o_acc+=o
            b_acc+=b
            i_acc+=i
            total_o+=num_o
            total_b+=num_b
            total_i+=num_i
            total += num
            prob = prob.numpy() #[B,38,3]

            #print(prob)
            with open("reslut"+str(epoch)+".txt",'a') as f:
                for i in range(prob.shape[0]):
                    for j in range(s_len[i]):
                        f.write('a \t b \t' +str(1-prob[i][j][0])+'\n')
                    f.write('\n')

    matchm,rankm = eval("result"+str(epoch)+".txt","train_dev_data/dev.txt")
    return matchm,rankm,sum(matchm.values()),acc/total, o_acc/total_o, b_acc/total_b, i_acc/total_i

def accuracy(prob,label,mask):
    #prob [B,38,3] label [B,38,3],mask:[B,38]

    prob_max = torch.argmax(prob,dim=2)
    label_max = torch.argmax(label,dim=2)
    acc = torch.sum((prob_max==label_max).float()*mask)
    o_prob = prob_max == 0
    o_label = label_max == 0
    b_prob = prob_max == 1
    b_label = label_max == 1
    i_prob = prob_max == 2
    i_label = label_max == 2
    o_acc = torch.sum((o_prob==o_label).float()*mask*o_label.float())
    b_acc = torch.sum((b_prob==b_label).float()*mask*b_label.float())
    i_acc = torch.sum((i_prob==i_label).float()*mask*i_label.float())
    total_o = torch.sum(o_label.float()*mask)
    total_b = torch.sum(b_label.float()*mask)
    total_i = torch.sum(i_label.float()*mask)
    return acc, o_acc, b_acc, i_acc, total_o,total_b,total_i,torch.sum(mask)

if __name__ == '__main__':
    args = parse_args()
    torch.manual_seed(args.seed)

    train_set = SEMDataset("train",'data/word2idx.pkl','data/pos2idx.pkl')
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)

    dev_set = SEMDataset('dev','data/word2idx.pkl','data/pos2idx.pkl')
    dev_loader = DataLoader(dev_set, batch_size=args.batch_size, shuffle=False)

    m = Model(args).to(device)

    #m = nn.DataParallel(m).to(device)
    train(m,train_loader,dev_loader,args)


import cv2
import mediapipe as mp
import math

video = cv2.VideoCapture('./estudos/cursos/udemy/Curso1/OpenCVBasic/ANEXO+polichinelos.mp4')
pose = mp.solutions.pose
Pose = pose.Pose(min_tracking_confidence=0.5, min_detection_confidence=0.5)
draw = mp.solutions.drawing_utils
count = 0
check = True

while True:
    check, img = video.read()

    videoRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    results = Pose.process(videoRGB)
    points = results.pose_landmarks
    draw.draw_landmarks(img, points, pose.POSE_CONNECTIONS)
    h, w, _ = img.shape

    if points:
        footRy = int(points.landmark[pose.PoseLandmark.RIGHT_FOOT_INDEX].y*h)
        footRx = int(points.landmark[pose.PoseLandmark.RIGHT_FOOT_INDEX].x*w)
        footLy = int(points.landmark[pose.PoseLandmark.LEFT_FOOT_INDEX].y*h)
        footLx = int(points.landmark[pose.PoseLandmark.LEFT_FOOT_INDEX].x*w)
        handRy = int(points.landmark[pose.PoseLandmark.RIGHT_INDEX].y*h)
        handRx = int(points.landmark[pose.PoseLandmark.RIGHT_INDEX].x*w)
        handLy = int(points.landmark[pose.PoseLandmark.LEFT_INDEX].y*h)
        handLx = int(points.landmark[pose.PoseLandmark.LEFT_INDEX].x*w)

        distHand = math.hypot(handRx-handLx, handRy-handLy)
        distFoot = math.hypot(footRx-footLx, footRy-footLy)

        if check == True and distHand <= 150 and distFoot >= 150:
            count += 1
            check = False
        if distHand > 150 and distFoot < 150:
            check = True

        text = f'Qtd: {count}'
        cv2.putText(img, text, (40,200), cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255))

    cv2.imshow('Video', img)
    cv2.waitKey(40)


import csv
import time
import random
import string
import numpy as np
from datetime import datetime
from datetime import timedelta
from dateutil.relativedelta import relativedelta

from ...helpers import get_file_lines, str_time_prop
from ...AbstractTableTransformer import AbstractTableTransformer

class ClientTableTransformer(AbstractTableTransformer):
    def __init__(self, settings, tableName):
        super().__init__(settings, tableName)

        self.clientsIdCounter = 0
        self.allClients = []

        # helpers
        self._date_format = self.settings["date_format"]
        self._date_format_out = self.settings["date_format_out"]

    def __get_ageGroup(self, birthDate):
        date = datetime.strptime(birthDate, self._date_format)
        age = relativedelta(datetime.now(), date).years

        age_group = "Unknown"
        if age <= 25:
            age_group = "15-25"
        elif age <= 35:
            age_group = "26-35"
        elif age <= 45:
            age_group = "36-45"
        elif age <= 55:
            age_group = "46-55"
        elif age <= 65:
            age_group = "56-65"
        elif age <= 75:
            age_group = "66-75"
        elif age <= 85:
            age_group = "76-85"
        elif age >= 86:
            age_group = "86-95"

        return age_group

    def __get_gender(self, gender):
        return gender

    def __get_married(self, married):
        status = "Unknown"

        if married == "married":
            status = "true"
        elif married == "single":
            status = "false"

        return status

    def __get_children(self, children):
        return children

    def __asserts(self, data):
        pass

    def __check_for_repeaters(self, data):
        for i in self.allClients:
            if data["AGE_GROUP"] == i["AGE_GROUP"] and data["GENDER"] == i["GENDER"] and data["MARRIED"] == i["MARRIED"] and data["CHILDREN"] == i["CHILDREN"]:
                data = i
                break
        else:
            data["CLIENTID"] = self.clientsIdCounter
            self.allClients.append(data)
            self.clientsIdCounter += 1
        return data

    def transform(self, src_data):
        # to collection
        data = {
            "CLIENTID": None,
            "AGE_GROUP": self.__get_ageGroup(src_data["client_birthDate"]),
            "GENDER": self.__get_gender(src_data["client_gender"]),
            "MARRIED": self.__get_married(src_data["client_married"]),
            "CHILDREN": self.__get_children(src_data["client_children"])
        }

        self.__asserts(data)
        data = self.__check_for_repeaters(data)
        self.__asserts(data)

        return data

    def table_write(self, data):
        self._writer.writerow([
            data["CLIENTID"],
            data["AGE_GROUP"],
            data["GENDER"],
            data["MARRIED"],
            data["CHILDREN"]
        ])

    def table_write_all(self):
        for item in self.allClients:
            self.table_write(item)


import sys
import matplotlib.pyplot as plt

def fib(n):
    fib.counter+=1
    if n == 0:
      return 0
    elif n == 1:
      return 1
    else:
      return fib(n-1) + fib(n-2)

fib.counter = 0
fib_mem.counter = 0

def plot(n, fun):
    x = []
    y = []
    for i in range(n):
        fun.counter = 0
        x.append(i)
        fun(i)
        y.append(fun.counter)
    plt.plot(x,y)
    plt.show()

def test(n):
    for i in range(n):
        assert(fib(i) == fib_mem(i))

if __name__ == '__main__':
    fib_mem(6)
    #test(18)
    #n = int(sys.argv[1])
    #plot(n, fib)
    #plot(n, fib_mem)


# Solution 1, Top-Down DP, memoization
class Solution:
    def dfs(self, i, j, A, B, d):
        # number of lines for the first i-th and j-th chars in A and B
        if (i, j) in d:
            return d[i, j]
        if i <= 0 or j <= 0:
            return 0

        if A[i - 1] == B[j - 1]:
            res = self.dfs(i - 1, j - 1, A, B, d) + 1
        else:
            res = max(self.dfs(i, j - 1, A, B, d), self.dfs(i - 1, j, A, B, d))
        d[i, j] = res

        return res

    def maxUncrossedLines(self, A: List[int], B: List[int]) -> int:
        m, n = map(len, (A, B))

        d = {}

        res = self.dfs(m, n, A, B, d)

        return res

# Solution 2, bottom-up DP
# notice this could be optimized to O(n) memory
class Solution:
    def maxUncrossedLines(self, A: List[int], B: List[int]) -> int:
        m, n = map(len, (A, B))

        dp = [[0 for _ in range(n + 1)] for _ in range(m + 1)]

        res = 0
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                if A[i - 1] == B[j - 1]:
                    dp[i][j] = 1 + dp[i - 1][j - 1]
                else:
                    dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
                res = max(res, dp[i][j])

        return res


# -*- coding: utf-8 -*-
"""
Catherine Lee
Basic Neural Network
"""
#works 7/2/18
#10 epochs: .971 accuracy!!!

#import
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

#load mnist data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

#variables for the NN
learning_rate = 0.5
epochs = 10
batch_size = 100

# declare the training data placeholders
# input x - for 28 x 28 pixels = 784
x = tf.placeholder(tf.float32, [None, 784]) #change this later depending on the image size
# now declare the output data placeholder - 10 digits
y = tf.placeholder(tf.float32, [None, 10])

#weight and bias variables for input and hidden layer

# now declare the weights connecting the input to the hidden layer
W1 = tf.Variable(tf.random_normal([784, 300], stddev=0.03), name='W1')
b1 = tf.Variable(tf.random_normal([300]), name='b1')
# and the weights connecting the hidden layer to the output layer
W2 = tf.Variable(tf.random_normal([300, 10], stddev=0.03), name='W2')
b2 = tf.Variable(tf.random_normal([10]), name='b2')

#activation function / calc output of hidden layer
hidden_out = tf.add(tf.matmul(x,W1),b1)
hidden_out = tf.nn.relu(hidden_out) #relu activation layer

#softmax activation for output layer
y_ = tf.nn.softmax(tf.add(tf.matmul(hidden_out,W2),b2))

#cross entropy cost function for the optimization/backpropagation to work on
y_clipped = tf.clip_by_value(y_, 1e-10, 0.9999999) #makes sure we never get log(0), aka parametrizes trainig
cross_entropy = -tf.reduce_mean(tf.reduce_sum(y * tf.log(y_clipped) #cross entropy calculation
                         + (1 - y) * tf.log(1 - y_clipped), axis=1))

#optimizer/ peforms gradient descent and backpropogation
#optimizer - makes training faster ?
optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate).minimize(cross_entropy)

#var init
init_op = tf.global_variables_initializer()

#accuracy assesment operation
correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# start the session
with tf.Session() as sess:
   # initialise the variables
   sess.run(init_op)

   total_batch = int(len(mnist.train.labels) / batch_size) #calc total num of batches we are running in each trianing epoch

   for epoch in range(epochs):
       avg_cost = 0 #to keep track of avg cross entropy cost for each epoch

       for i in range(total_batch):
           batch_x, batch_y = mnist.train.next_batch(batch_size=batch_size) #get random batch x and batch y from MNIST
           _, c = sess.run([optimizer, cross_entropy],
                           feed_dict={x: batch_x, y: batch_y})
           avg_cost += c / total_batch
           print("Epoch:", (epoch + 1), "cost =", "{:.3f}".format(avg_cost))
   print(sess.run(accuracy, feed_dict={x: mnist.test.images, y: mnist.test.labels}))


import mne
import os
import glob
import pandas as pd
import gspread
from google.oauth2 import service_account
from datetime import datetime, timedelta
import shutil
import pytz
import numpy as np

# Load credentials and connect to Google Sheets API
scope = ['https://www.googleapis.com/auth/spreadsheets']
creds = service_account.Credentials.from_service_account_file(
    '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/GoogleSheet_api_key/api-access-381407-51bd2bb2d008.json', scopes=scope)
client = gspread.authorize(creds)

# Replace this with your Google Sheet key
with open('/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/GoogleSheet_api_key/sheet_key.txt') as f:
    contents = f.readlines()
sheet_key = contents[0]
sheet_name = 'Devices/Session Tracking'
# Read the data from the Google Sheet
sheet = client.open_by_key(sheet_key).worksheet(sheet_name)

data = sheet.get_all_records()

# Convert the data to a pandas DataFrame

df1 = pd.DataFrame(data)

# add subject 3101 to the df structure
sheet_name2 = 'NUS Onsite Dry-Run Data'
# Read the data from the Google Sheet
sheet2 = client.open_by_key(sheet_key).worksheet(sheet_name2)
data2 = sheet2.get_all_records()

df2 = pd.DataFrame(data2)

df = pd.concat([df1, df2])

data_dir = '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Zepplife/Zepp_Life_processed.csv'
HPB_BedwakeTime = pd.read_csv(data_dir)
subjects = (HPB_BedwakeTime['NUS_ID']).unique()

psg = '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Sleep_Staging/Consensus_score/'

Saving_file_path = "/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/Xiaomi"

# staging_file_path = "/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Sleep_Staging"
# Now you can test your for loop using the `data` DataFrame

combined_data_AllDays = pd.DataFrame()
combined_data_AllDays_diff = pd.DataFrame()
i = -1

for subj in subjects:
    i = i+1

    match = df[(df['Night'] == int(subj[-1]))
               & (df['Participant ID'] == subj[:-3])]

    lightOff = match.iloc[0]['Session start local time']
    lightON = match.iloc[0]['Session end local time']

    lightOff = pd.Timestamp(lightOff)
    base_date = pd.Timestamp(match.iloc[0]['Session end date'])
    if lightOff.hour > 19:
        lightOff_day = base_date - pd.Timedelta(days=1)
        lightOff = lightOff.replace(year=lightOff_day.year,
                                    month=lightOff_day.month, day=lightOff_day.day).tz_localize(pytz.UTC)
    else:
        lightOff = lightOff.replace(year=base_date.year,
                                    month=base_date.month, day=base_date.day).tz_localize(pytz.UTC)
    lightON = pd.Timestamp(lightON).replace(
        year=base_date.year, month=base_date.month, day=base_date.day).tz_localize(pytz.UTC)

    HPB_BedwakeTime_tmp= HPB_BedwakeTime[HPB_BedwakeTime['NUS_ID']== subj]

    subject_psg=subj.replace('NUS3_','NUS3')
    psg_file = f'{psg}{subject_psg}*sleepstage_consensus.csv'
    PSG_data = glob.glob(psg_file)
    df_PSG = pd.read_csv(PSG_data[0])

    # first_sleep_ep = df_PSG.ne(0).idxmax().to_frame('pos')
    # last_sleep_ep = df_PSG.iloc[::-1, :].ne(0).idxmax().to_frame('pos')

    first_sleep_ep = df_PSG.ne(0).idxmax()
    last_sleep_ep = df_PSG.iloc[::-1, :].ne(0).idxmax()

    A = pd.to_datetime(
        HPB_BedwakeTime_tmp['start_stage'].iloc[0], utc=True)

    B = lightOff+pd.Timedelta(minutes= int(first_sleep_ep)/2)

    hour = int(A.strftime('%H'))
    minutes = int(A.strftime('%M'))
    seconds = int(A.strftime('%S'))
    hour_in_seconds_A1 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_A1 = hour_in_seconds_A1+hr

    hour = int(B.strftime('%H'))
    minutes = int(B.strftime('%M'))
    seconds = int(B.strftime('%S'))
    hour_in_seconds_B1 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_B1 = hour_in_seconds_B1+hr

    diff = hour_in_seconds_A1-hour_in_seconds_B1
    print(
        f"Bedtime start difference Based on Seconds  {subj}  NIGHT 1    {diff} minutes            OURA {A }  and  PSG {B}")

    # Calculate the time difference and include the sign
    time_diff = A - B
    diff_seconds1 = time_diff.total_seconds()

    # difference of Lights ON and Oura start of recording
    BedtimeStart_difference = diff_seconds1

    print(
        f"Bedtime start difference  {subj}  NIGHT 1   {diff_seconds1/60} minutes            OURA {A }  and  PSG {B}")

    # A2 = df_Oura['timestamp'][len(df_Oura)-1]

    A = pd.to_datetime(
        HPB_BedwakeTime_tmp['end_stage'].iloc[-1], utc=True)
    B = lightON-pd.Timedelta(minutes= (len(df_PSG)-int(last_sleep_ep))/2)

    hour = int(A.strftime('%H'))
    minutes = int(A.strftime('%M'))
    seconds = int(A.strftime('%S'))
    hour_in_seconds_A2 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_A2 = hour_in_seconds_A2+hr

    hour = int(B.strftime('%H'))
    minutes = int(B.strftime('%M'))
    seconds = int(B.strftime('%S'))
    hour_in_seconds_B2 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_B2 = hour_in_seconds_B2+hr

    diff = hour_in_seconds_A2-hour_in_seconds_B2

    # Calculate the time difference and include the sign
    time_diff2 = A - B
    diff_seconds2 = time_diff2.total_seconds()

    BedtimeEnd_difference = diff_seconds2
    print(
        f"Waketime  difference    {subj}  NIGHT 1    {BedtimeEnd_difference} minutes            OURA {A }  and  PSG {B}")

    start_ep_diff = BedtimeStart_difference/30
    end_ep_diff = BedtimeEnd_difference/30

    combined_data_diff = pd.DataFrame({
        'subj': [subj],

        'StartTime_device': [hour_in_seconds_A1],
        'StartTime_ref': [hour_in_seconds_B1],
        'EndTime_device': [hour_in_seconds_A2],
        'EndTime_ref': [hour_in_seconds_B2]

    })
    if int(subj[-1])==1 :
        combined_data_AllDays_diff = combined_data_AllDays_diff.append(
            combined_data_diff, ignore_index=True
        )
        print(combined_data_AllDays_diff)

combined_data_AllDays_diff.to_csv(os.path.join(Saving_file_path,
                                               f"SOL_waketime_Diff_data_N1.csv"), index=False)

################################################################
################################################################
################################################################


CAPACITY = 10
class Heap:
    def __init__(self):
        self.heap = [0]*CAPACITY
        self.size = 0

    #use this method to insert a data
    def insert(self,data):
        if self.size == CAPACITY:
            return
        self.heap[self.size] = data
        self.size =self.size+1
        self.fix_heap(self.size-1)

    # arrange heap acc to property
    def fix_heap(self,index):
        parent_index = (index - 1)//2
        if index > 0 and self.heap[index] > self.heap[parent_index]:
            self.swap(index,parent_index)
            self.fix_heap(parent_index)

    def swap(self,index1,index2):
        self.heap[index2],self.heap[index1] = self.heap[index1],self.heap[index2]

    def getmax(self):
        return self.heap[0]

    def poll(self):
        max = self.getmax()
        self.swap(0,self.size-1)
        #del self.heap[self.size -1]
        self.size = self.size-1
        self.fix_down(0)
        return max
    # to remove violations in downward direction
    def fix_down(self,index):
        left_index = 2*index + 1
        right_index = 2*index + 2

        large_value_index = index
        if left_index < self.size and self.heap[large_value_index] < self.heap[left_index]:
            large_value_index = left_index
        if right_index < self.size and self.heap[large_value_index] < self.heap[right_index]:
            large_value_index = right_index
        if index != large_value_index:
            self.swap(index,large_value_index)
            self.fix_down(large_value_index)

    def heap_sort(self):
        size = self.size
        for i in range(0,size):
            max =self.poll()
            print (max)

if __name__ == "__main__":
    Heap = Heap()
    Heap.insert(10)
    Heap.insert(12)
    Heap.insert(1)
    Heap.insert(19)
    Heap.insert(0)

    Heap.heap_sort()


from app import models, test
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from .database import engine

models.Base.metadata.create_all(bind=engine)
app = FastAPI()

origins = [
    "*"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)
app.include_router(test.router, tags=['Tests'], prefix='/api/tests')
@app.get("/api")
def root():
    return {"message": "welcome to this test fast api server"}


from api.file import File

file = File()
file.filename = "scenario1"
file.title = "test.txt"
file.contents = "Dette er en test"

file.new_key("testkey")

file.config['path'] =  f'{file.root_directory}/scenarios/'
file.create_directories()
file.save()


import numpy as np
from sklearn import svm
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier, GradientBoostingClassifier
from sklearn.tree import DecisionTreeClassifier

def get_labeled_windowed_data(observations, window_size=7):
    """
    Split up the observations into windowed chunks. Each windowed chunk of
    observations is associated with a label vector of what the price change is
    per market *immediately after* the windowed chunk (+1 for price goes up,
    0 for no change, and -1 for price goes down). Thus, a classifier's task
    for the data is given a windowed chunk, to predict what its label is
    (i.e., given recent percent changes in all the markets, predict the
    directions of the next price changes per market).

    Inputs
    ------
    - observations: 2D array; each column is a percent-change time series data
        for a specific market
    - window_size: how large the window is (in number of time points)

    Outputs
    -------
    - windows: 3D array; each element of the outermost array is a 2D array
        of the same format as `observations` except where the number of time
        points is exactly `window_size`
    - window_labels: 2D array; `window_labels[i]` is a 1D vector of labels
        corresponding to the time point *after* the window specified by
        `windows[i]`; `window_labels[i]` says what the price change is for
        each market (+1 for going up, 0 for staying the same, and -1 for going
        down)

    *WARNING*: Note that the training data produced here is inherently not
    i.i.d. in that `windows[0]` and `windows[1]`, for instance, will largely
    overlap!
    """
    num_time_points, num_markets = observations.shape
    windows = []
    window_labels = []
    for start_idx in range(num_time_points-window_size):
        windows.append(observations[start_idx:start_idx+window_size])
        window_labels.append(1*(observations[start_idx+window_size] > 0)
                             -1*(observations[start_idx+window_size] < 0))
    windows = np.array(windows)
    window_labels = np.array(window_labels)
    return windows, window_labels

# global variables to be saved for the trained classifier
clfs = None
clfs1 = None

def train(windows, window_labels):
    """
    Your training procedure goes here! It should train a classifier where you
    store whatever you want to store for the trained classifier as *global*
    variables. `train` will get called exactly once on the exact same training
    data you have access to. However, you will not get access to the mystery
    test data.

    Inputs
    ------
    - windows, window_labels: see the documentation for the output of
        `get_labeled_windowed_data`
    """

    # -------------------------------------------------------------------------
    # YOUR CODE HERE
    #

    # The autograder wants you to explicitly state which variables are global
    # and are supposed to thus be saved after training for use with prediction.
    global clfs, clfs1, clfs2

    #guess = 0

    #from sklearn.metrics import log_loss

    n, m = 7, 5
    X = np.concatenate((windows[:,-n:,0], windows[:,-n:,1], windows[:,-n:,2], windows[:,-n:,3]), axis=1)
    for i in range(4):
        for k in range(2,8):
            X = np.concatenate((X, ((windows[:,-k:,i] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            X = np.concatenate((X, ((windows[:,-k:,i] > 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            X = np.concatenate((X, ((windows[:,-k:,i] < 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,1] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,2] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,3] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] == 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] > 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        #X = np.concatenate((X, ((windows[:,-m:,0] >= 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] < 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)

    clfs = [RandomForestClassifier(n_estimators=10) for i in range(4)]
    #clfs = [ AdaBoostClassifier(DecisionTreeClassifier(max_depth=1),
    #                     algorithm="SAMME",
    #                    n_estimators=20) for i in range(4)]
    clfs1 = [GradientBoostingClassifier(n_estimators=100, learning_rate=1.0, max_depth=1, random_state=0) for i in range(4)]
    clfs2 = [svm.SVC(C=10**6, kernel='rbf') for i in range(4)]

    for i in range(4):
       clfs[i] = clfs[i].fit(X, window_labels[:,i])
       clfs1[i] = clfs1[i].fit(X, window_labels[:,i])
       clfs2[i] = clfs2[i].fit(X, window_labels[:,i])
       #clf_probs = clf.predict_proba(X)
       #score = log_loss(y, clf_probs)
    #
    # END OF YOUR CODE
    # -------------------------------------------------------------------------

def forecast(window):
    """
    Your forecasting method goes here! You may assume that `train` has already
    been called on training data and so any global variables you stored as a
    result of running `train` are available to you here for prediction
    purposes.

    Input
    -----
    - window: 2D array; each column is 7 days worth of percent changes in
        price for a specific market

    Output
    ------
    1D array; the i-th entry is a prediction for whether the percentage
    return will go up (+1), stay the same (0), or go down (-1) for the i-th
    market
    """

    # -------------------------------------------------------------------------
    # YOUR CODE HERE
    #
    #from collections import Counter

    global clfs, clfs1, clfs2
    from scipy.stats import mode

    n, m = 7, 5
    X = window[-n:,:]
    X = X.flatten('F')
    for i in range(4):
        for k in range(2,8):
            X = np.append(X, (window[-k:,i] == 0).sum()==k)
            X = np.append(X, (window[-k:,i] > 0).sum()==k)
            X = np.append(X, (window[-k:,i] < 0).sum()==k)
            #X = np.append(X, (window[-k:,1] == 0).sum()==k)
            #X = np.append(X, (window[-k:,2] == 0).sum()==k)
            #X = np.append(X, (window[-k:,3] == 0).sum()==k)
        X = np.append(X, (window[-m:,i] == 0).sum())
        X = np.append(X, (window[-m:,i] > 0).sum())
        #X = np.append(X, (window[-m:,0] >= 0).sum())
        X = np.append(X, (window[-m:,i] < 0).sum())

    predicted_labels = np.array([mode([clfs[idx].predict(X), clfs1[idx].predict(X), clfs2[idx].predict(X)], axis=0).mode.flatten() for idx in range(window.shape[1])])

    #
    # END OF YOUR CODE
    # -------------------------------------------------------------------------

    return predicted_labels

def main():
    # get coconut oil challenge training data
    observations = []
    with open('coconut_challenge.csv', 'r') as f:
        for line in f.readlines():
            pieces = line.split(',')
            if len(pieces) == 5:
                observations.append([float(pieces[1]),
                                     float(pieces[2]),
                                     float(pieces[3]),
                                     float(pieces[4])])
    observations = np.array(observations)
    train_windows, train_window_labels = \
        get_labeled_windowed_data(observations, window_size=7)

    train(train_windows, train_window_labels)

    # figure out accuracy of the trained classifier on predicting labels for
    # the training data
    train_predictions = []
    for window, window_label in zip(train_windows, train_window_labels):
        train_predictions.append(forecast(window))
    train_predictions = np.array(train_predictions)

    train_prediction_accuracy_plus1 = \
        np.mean(train_predictions[train_window_labels == 1]
                == train_window_labels[train_window_labels == 1])
    train_prediction_accuracy_minus1 = \
        np.mean(train_predictions[train_window_labels == -1]
                == train_window_labels[train_window_labels == -1])
    train_prediction_accuracy_0 = \
        np.mean(train_predictions[train_window_labels == 0]
                == train_window_labels[train_window_labels == 0])
    print('Training accuracy for prediction +1:',
          train_prediction_accuracy_plus1)
    print('Training accuracy for prediction -1:',
          train_prediction_accuracy_minus1)
    print('Training accuracy for prediction 0:',
          train_prediction_accuracy_0)

if __name__ == '__main__':
    main()

courses/MITx/6.008.1x - Computational Probability & Inference (2017)/Week11/final_proj_challenge.py


import os
import time
import pandas as pd
import numpy as np
from bifac import BiFac
from sklearn.metrics import normalized_mutual_info_score
from sklearn.metrics import f1_score

#preparation
logdir = "log/logdaisharin"
os.system("rm -rf log/logdaisharin")
df = pd.read_csv("data/edge_list.csv", header=None)
edge_list = [(i-1, j-1) for i, j in zip(df[0], df[1])]
weights = df[2]
ans = pd.read_csv("data/term_category.csv", header=None)[1]
N = ans.shape[0]
cnst = []
ncnst = 100
while len(cnst) < ncnst:
    n1 = np.random.randint(0, N)
    n2 = np.random.randint(0, N)
    if n1 == n2: continue
    if ans.ix[n1] == ans.ix[n2]:
        cnst.append((n1,n2))

#construct BiFac model
model = BiFac(K=8, edge_list=edge_list, weights=weights, learning_rate=1.0,
              lambda_r=0, constraint=cnst, lambda_c=0.)

#detect communities
start = time.time()
loss, z, theta1, theta2 = model.optimize(logdir=logdir, stop_threshold=10e-8,
                                        max_steps=5000)
sumtime = time.time() - start

com1, com2 = model.get_hard_community()

fvalue = f1_score(com1, ans)
nmi = normalized_mutual_info_score(com1, ans)

print("time", sumtime)
print("fvalue", fvalue)
print("nmi", nmi)


import pandas as pd
from mlxtend.frequent_patterns import apriori,association_rules

data=pd.read_csv("D:\\Data Science\\Assignment files\\Association Rules Assignment\\transactions_retail1.csv",header=None)
num_records=len(data)
data.describe()

#Imputation of NULL values#
data.isnull().sum()
data.shape
data.dropna().shape

data[1].fillna("HEART",inplace=True)
data[5].fillna("SET",inplace=True)
data[4].fillna("SET",inplace=True)
data[3].fillna("OF",inplace=True)
data[2].fillna("RED",inplace=True)

records=[]

for i in range(557):

    records.append([str(data.values[i,j]) for j in range (0,5)])

all_data_list= [i for item in records for i in item]

from collections import Counter
item_frequencies=Counter(all_data_list)

item_frequencies=sorted(item_frequencies.items(),key=lambda x:x[1])

records_series=pd.DataFrame(pd.Series(records))
records_series.columns=["transactions"]

# creating a dummy columns for the each item in each transactions ... Using column names as item name
X = records_series['transactions'].str.join(sep='*').str.get_dummies(sep='*')

frequent_itemsets = apriori(X,min_support=0.005, max_len=3,use_colnames = True)
frequent_itemsets.shape

# Most Frequent item sets based on support
frequent_itemsets.sort_values('support',ascending = False,inplace=True)
rules = association_rules(frequent_itemsets, metric="lift", min_threshold=1)
rules.shape

rules.head(10)
rules.sort_values('lift',ascending = False,inplace=True)


import torch
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import argparse
import torchvision.models as models
from classification_dataset import DVTDataset
from torch import nn
from sklearn.metrics import roc_auc_score, precision_score, recall_score
import numpy as np
import re
from torch.utils.tensorboard import SummaryWriter
import json

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

def evaluate(model, dataloader, criterion, writer=None, epoch=0, val=False):
    auc = 0
    prec = 0
    recall = 0
    model = model.eval()
    total_preds = np.array([])
    total_labels = np.array([])
    total_hard_preds = np.array([])

    for i, data in enumerate(dataloader):
        inputs, labels, access_number = data['image'], data['label'], data['access_number']
        inputs = inputs.repeat(1, 3, 1, 1).float()
        inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        if val:
            # print(torch.argmax(outputs, 1))
            loss = criterion(outputs, labels)
            preds = torch.argmax(outputs, 1)
            # print(preds)
            # auc = auc + roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy())
            # prec = prec + precision_score(labels.cpu().numpy(), preds.cpu().numpy())
            # recall = recall + recall_score(labels.cpu().numpy(), preds.cpu().numpy())
            # print("AUC: ", roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("prec: ", precision_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("recall: ", recall_score(labels.cpu().numpy(), preds.cpu().numpy()))
            if i % args.log_interval == 0:
                batch_number = i + epoch * len(dataloader)
                print("Val Batch Number: {} | Val Loss: {} ".format(batch_number, loss))
                writer.add_scalar("Validation Loss", loss, batch_number)

        else:
            # Compute AUC/ Precision
            preds = torch.argmax(outputs, 1)
            # soft = nn.Softmax()
            # out = soft(outputs)
            # print(preds)
            try:
                # print("Labels", labels)
                # print("Preds", preds)
                out_preds = torch.max(outputs, 1)[0]
                out_hard_preds = torch.argmax(outputs, 1)
                total_hard_preds = np.concatenate((total_hard_preds, out_hard_preds.cpu().detach().numpy()))
                total_preds = np.concatenate((total_preds, out_preds.cpu().detach().numpy()))
                total_labels = np.concatenate((total_labels, labels.cpu().numpy()))
                current_prec = precision_score(labels.cpu().numpy(), preds.cpu().numpy())
                current_recall = recall_score(labels.cpu().numpy(), preds.cpu().numpy())
                # print("Current Prec", current_prec)
                # print("Current Recall", current_recall)
                # print("Equal", current_prec == 1.0)
                auc = auc + roc_auc_score(labels.cpu().numpy(), preds.cpu().detach().numpy())
                prec = prec + precision_score(labels.cpu().numpy(), preds.cpu().numpy())
                recall = recall + recall_score(labels.cpu().numpy(), preds.cpu().numpy())

            except ValueError:
                pass
            # print("AUC: ", roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("prec: ", precision_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("recall: ", recall_score(labels.cpu().numpy(), preds.cpu().numpy()))

    auc = float(auc) / i
    prec = float(prec) / i
    recall = float(recall) / i
    print("Average AUC: ", auc)
    print("Average Prec: ", prec)
    print("Average Recall: ", recall)
    d = {'preds': list(total_preds), 'labels': list(total_labels), 'hard_preds': list(total_hard_preds)}
    jn = json.dumps(d)
    f = open("dict_2.json", "w")
    f.write(jn)
    f.close()
    # writer.add_scalar("Validation Loss", loss, i)
    # writer.add_scalar("Validation AUC", auc, i)
    # writer.add_scalar("Validation Precision", prec, i)
    # writer.add_scalar("Validation Recall", recall, i)
    return auc

def main(args):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    train_dataset = DVTDataset(csv_file=args.train, root_dir=args.root_dir, device=device)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)

    model = models.densenet161()
    print("Loading model onto CUDA")
    state_dict = torch.load("densenet161-8d451a50.pth")
    pattern = re.compile(
        r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')

    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]
    model.load_state_dict(state_dict)
    model.classifier = nn.Linear(2208, 2)

    model = model.cuda()
    model.train()

    print("Finished loading model onto CUDA")

    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    # optimizer = optim.Adam(model.parameters(), lr = 0.0001)
    # optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)
    # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
    writer = SummaryWriter("classification_lowlr_runs")
    # weights = np.zeros(1000)
    # weights[0] = 1.0
    # weights[1] = 1.0
    criterion = nn.CrossEntropyLoss()
    if not args.evaluate:
        for epoch in range(args.epochs):
            for i, data in enumerate(train_loader):
                # get the inputs; data is a list of [inputs, labels]
                inputs, labels = data['image'], data['label']
                # print(inputs.shape)
                inputs = inputs.repeat(1, 3, 1, 1).float()
                inputs, labels = inputs.cuda(), labels.cuda()
                # print(inputs.shape)
                # zero the parameter gradients
                optimizer.zero_grad()
                # print("Forward pass")
                # forward + backward + optimize

                outputs = model(inputs)
                # print(outputs)
                # print(torch.argmax(outputs, 1))
                # print("Forward pass complete.")
                loss = criterion(outputs, labels)
                loss.backward()
                # print("Backward pass.")
                optimizer.step()
                if i % args.log_interval == 0:
                    batch_number = i + epoch * len(train_loader)
                    # print(torch.argmax(outputs, 1))
                    writer.add_scalar("Training Loss", loss, batch_number)
                    print("Batch Number: {} | Loss: {}".format(i, loss))
            torch.save(model.state_dict(), "classification_lowlr/model_{}.pt".format(epoch))
            val_dataset = DVTDataset(csv_file=args.val, root_dir=args.root_dir, device=device)
            val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True,
                                                     num_workers=0)
            evaluate(model, val_loader, criterion, writer, epoch, val=True)
            # scheduler.step()
            # test_dataset = DVTDataset(csv_file = args.test, root_dir = args.root_dir, device = device)
            # test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = args.batch_size, shuffle = True, num_workers = 0)
            # evaluate(model, test_loader, criterion)

    model.load_state_dict(torch.load(args.checkpoint_name + "classification_dropout/model_9.pt"))
    model.eval()
    test_dataset = DVTDataset(csv_file=args.test, root_dir=args.root_dir, device=device)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
    evaluate(model.eval(), test_loader, criterion)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', default='train_femoral.csv')
    parser.add_argument('--val', default='validation_femoral.csv')
    parser.add_argument('--test', default='train_femoral.csv')
    parser.add_argument('--root_dir', default='')
    parser.add_argument('--batch_size', default=10)
    parser.add_argument('--epochs', default=10)
    parser.add_argument('--lr', default=0.001)
    parser.add_argument('--log_interval', default=10)
    parser.add_argument('--val_interval', default=100)
    parser.add_argument('--evaluate', default=False)
    parser.add_argument('--checkpoint_name', default="")

    args = parser.parse_args()
    main(args)


#!/usr/bin/env python3.8

## LIBRARIES
import argparse
from Bio import SeqIO
import sys
import time
import datetime
import glob
import os
from PDOR_lib import core_snps_2_fasta
from PDOR_lib import core_snps_list_path

## FUNCTIONS
class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

def parse_args():

    parser=MyParser(usage='%(prog)s -q QUERY_FOLDER -sd SD_SKETCH -ref REFERENCE_GENOME -snp_thr SNP_THRESHOLD [options]',
        prog='P-DOR.py',
        argument_default=argparse.SUPPRESS,
        epilog="According to the legend, P-dor is the Son of K-mer, but it also likes SNPs")
    requiredNamed = parser.add_argument_group('Input data (all required, except when in TEST mode)')
    requiredNamed.add_argument('-q', help='query folder containing genomes in .fna format',metavar="<dirname>",dest="query_folder",type=str,required=False,default=None)
    requiredNamed.add_argument("-sd", help="Source Dataset (SD) sketch file",metavar="<dirname>",dest="db_sketch", required=False,default=None)
    requiredNamed.add_argument("-ref", help="reference genome", dest="ref", metavar="<filename>",required=False,default=None)
    requiredNamed.add_argument("-snp_thr", dest="snp_threshold", help="Threshold number of SNPs to define an epidemic cluster",required=False,metavar="<int>",default=None)

    optional = parser.add_argument_group('Additional arguments')
    optional.add_argument("-amrf", help="if selected, P-DOR uses amrfinder-plus to search for antimicrobial resistance and virulence genes in the entire Analysis Dataset [Default: %(default)s]",action="store_true",required=False)
    optional.add_argument("-meta", help="metadata file; see example file for formatting [Default: %(default)s]",default=None,required=False)
    optional.add_argument("-sd_folder", help="folder containing the genomes from which the Source Dataset (SD) sketch was created [Default: P-DOR downloads the background genomes from BV-BRC]",required=False,default="",dest="bkg_folder",nargs="?")
    optional.add_argument("-borders", type=int, help="length of the regions at the contig extremities from which SNPs are not called [Default: %(default)s]",metavar="<int>",default=20)
    optional.add_argument("-min_contig_length", type=int, help="minimum contig length to be retained for SNPs analysis [Default: %(default)s]",metavar="<int>",default=500)
    optional.add_argument("-snp_spacing", type=int, help="Number of bases surrounding the mutated position to call a SNP [Default: %(default)s]",metavar="<int>",default=10)
    optional.add_argument('-species', help="full species name of the genomes in the query folder [Default: %(default)s]\nNeeds to be written within single quotes, e.g.: -species 'Klebsiella pneumoniae'.\nThis option is used for a quicker and more precise gene search with AMRFinderPlus", default="")
    optional.add_argument('-n', type=int,help='Maximum closest genomes from Source Dataset (SD) [Default: %(default)s]',metavar="<int>",dest="near",default=20)
    optional.add_argument('-t', type=int,help='number of threads [Default: %(default)s]',metavar="<int>",dest="threads",default=2)
    parser.add_argument('-TEST', action= 'store_true',help='TEST MODE. Warning: this option overrides all other arguments', required=False)
    parser.add_argument('-v', '--version', action='version', version='P-DOR v1.1')

    parser.set_defaults(amrf=False)
    parser.set_defaults(TEST=False)
    args = parser.parse_args()
    if len(sys.argv)==1:
        parser.print_help()
        sys.exit()
    elif (not args.TEST and (args.query_folder is None or args.db_sketch is None or args.ref is None or args.snp_threshold is None)):
        parser.error("\n\nERROR!! If not in TEST mode, the following arguments are required: -q, -sd, -ref, -snp_thr\n\n")
    return args

def logfile():
	with open("PDOR.log","w") as p:
		for arg, value in sorted(vars(args).items()):
			p.write("%s\t%s\n" %(arg,value))

def parse_species(sp):
    orglist=["Acinetobacter_baumannii", "Burkholderia_cepacia", "Burkholderia_pseudomallei", \
    "Campylobacter", "Clostridioides_difficile", "Enterococcus_faecalis", "Enterococcus_faecium", \
    "Escherichia", "Klebsiella_oxytoca", "Klebsiella_pneumoniae", "Neisseria_gonorrhoeae", \
    "Neisseria_meningitidis", "Pseudomonas_aeruginosa", "Salmonella", "Staphylococcus_aureus", \
    "Staphylococcus_pseudintermedius", "Streptococcus_agalactiae", "Streptococcus_pneumoniae", \
    "Streptococcus_pyogenes", "Vibrio_cholerae"]
    organism=""
    if sp.strip()!="":
        test="_".join(sp.strip().split())
        if test in orglist:
            organism=test
        else:
            test=sp.strip().split()[0]
            if test in orglist:
                organism=test
    return organism

def check_res_vir(threads,AD_folder,i_ref):
	print ("Checking for resistance and virulence genes...")
	path_res=path_dir+"/"+Results_folder_name+"/"+AD_folder
	os.chdir(path_res)
	os.system("cp ../%s ." %(i_ref))
	os.system("amrfinder --update")
	organism=parse_species(args.species)
	if organism=="":
		os.system("ls *fna | parallel -j %i 'amrfinder -n {} --plus --threads 1 -o {}.txt --name {}'" %(threads))
	else:
		os.system("ls *fna | parallel -j %i 'amrfinder -n {} --plus --threads 1 -o {}.txt -O %s --name {}'" %(threads,organism))
	os.system("cat *fna.txt >summary_resistance_virulence.txt")
	os.system("rm *fna.txt")
	os.system("rm %s/%s" %(path_res,i_ref.split("/")[-1]))

	os.system("mv summary_resistance_virulence.txt %s/%s" %(path_dir,Results_folder_name))
	os.chdir("..")

def Mummer_snp_call(threads,i_ref,Align_folder):
	print("Aligning genomes with Mummer4...\n")
	os.chdir(path_dir+"/"+Results_folder_name+"/"+Align_folder)
	os.system('ls *fna | parallel -j %i "nucmer --maxgap=500 -p {} %s {}; delta-filter -1 {}.delta > {}_filtered.delta; show-snps -H -C -I -r -T {}_filtered.delta | cut -f1,2,3 >{}.snp"' %(threads,i_ref))
	os.system('ls *_filtered.delta | parallel -j %i "dnadiff -d {} -p {}_info >/dev/null 2>&1"' %(threads))
	reportFiles=glob.glob("*.report")
	rFdict={}
	for rF in reportFiles:
		rFcontent=open(rF,"r").read().split("\n")
		rFdict[rF.strip().replace(".fna_filtered.delta_info.report","")]=float([t.strip().split()[-1].split("(")[1].split("%")[0] for t in rFcontent if "AlignedBase" in t][0])
	with open("Coverage_report.txt","w") as rFout:
		sleep=False
		for t in rFdict.keys():
			if rFdict[t]<70:
				sleep=True
				print("\n\nWARNING! Genome %s aligns poorly to the REFERENCE GENOME (%f %% horizontal coverage). Please, check!\n\n" %(t,rFdict[t]))
			rFout.write("%s\t%f\n" %(t,rFdict[t]))
		if sleep==True:
			time.sleep(10)
	#sys.exit()
	os.chdir(path_dir+"/"+Results_folder_name)
	os.system("ls %s/*.snp >SNP_genomes.list" %(Align_folder))
	core_snps_list_path("SNP_genomes.list", snp_spacing, "SNP_positions")
	core_snps_2_fasta("SNP_genomes.list", "SNP_positions.core.list",i_ref,"SNP_alignment")
	if os.stat("SNP_alignment.core.fasta").st_size > 10:
		print ("Core-SNPs aligmment detected!\n")
	else:
                sys.exit("\nERROR: Core-SNPs alignment not present, check your input files!\n")

### MAIN
print('\n\n      (/T\)  _____\n      (-,-) ()____)\n      \(o)/   -|3       BEHOLD! This is P-DOR!\n     /="""===//||\n,OOO//|   |    ""\nO:O:O LLLLL\n\OOO/ || ||\n     C_) (_D\n\n')
args = parse_args()
time_now=datetime.datetime.now()
datestamp=time_now.strftime("%Y-%m-%d_%H-%M-%S")
start_time = time.time()
prog_dir = os.path.abspath( os.path.dirname( sys.argv[0]))
Results_folder_name="Results_%s" %(datestamp)

##CHECK IF TEST MODE
if args.TEST==True:
    Results_folder_name="Test_%s" %(datestamp)
    args.amrf=False
    args.bkg_folder='%s/data/test_DB/' %(prog_dir)
    args.borders=20
    args.db_sketch='%s/data/sketches.msh' %(prog_dir)
    args.meta='%s/data/sample_metadata_table.txt' %(prog_dir)
    args.min_contig_length=500
    args.near=5
    args.query_folder='%s/data/test_query/' %(prog_dir)
    args.ref='%s/data/NJST258.fna' %(prog_dir)
    args.snp_spacing=10
    args.snp_threshold='20'
    args.species='Klebsiella pneumoniae'
    args.threads=2
    print("\n\nTesting AMRFinderPlus function on a single genome")
    os.system("amrfinder --update")
    organism=parse_species(args.species)
    os.system("amrfinder -n %s --plus --threads %i -o test.amrftest -O %s --name REF" %(args.ref,args.threads,organism))
    os.system("head -n2 test.amrftest")
    os.system("rm test.amrftest")

os.mkdir(Results_folder_name)
logfile()
path_dir = os.path.abspath( os.path.dirname( Results_folder_name))

os.system("mv PDOR.log %s" %(Results_folder_name))

query_folder=args.query_folder
db_sketch=os.path.abspath(args.db_sketch)
ref=args.ref
threads=args.threads
nearest=args.near
borders=args.borders
contig_length=args.min_contig_length
snp_spacing=args.snp_spacing
bkg_folder=args.bkg_folder
if args.meta is not None:
    ABSmeta=os.path.abspath(args.meta)

ref=os.path.abspath(ref)
ABS_query_folder=os.path.abspath(query_folder)

exts=['*.fasta', '*.fna', '*.fa']

files = [f for ext in exts for f in glob.glob(os.path.join(ABS_query_folder, ext))]

genomes_query = [i.strip().split("/")[-1] for i in files]
#genomes_query=[i.replace(".fasta",".fna").replace(".fa",".fna") for i in genomes_query]

print ("\nChecking input formats...\n")
### CHECK N 1
# REF has 1 contig
num_cont_c1=int(os.popen('grep -c ">" %s' %(ref)).read().strip())
if num_cont_c1>1:
	sys.exit('\nERROR: the reference file needs to contain only one sequence!')
### CHECK N 2
# REF is in fasta format
elif num_cont_c1==0:
	sys.exit('\nERROR: the reference file needs to be a FASTA file!')
### CHECK N 3
if db_sketch.split(".")[-1]!="msh":
	sys.exit("\n\nERROR: please check you sketch file!!!")
### CHECK N 4
for gen in genomes_query:
	if gen.endswith(".fasta") or gen.endswith(".fna") or gen.endswith(".fa"):
		n_genomes=len(genomes_query)
		print ("Detected %i genomes in the query folder\n" %n_genomes)
		break
	else:
		sys.exit("\n\nERROR: your query genome folder does not contain fasta files!!!")

### format REF
alt_ref=SeqIO.read(ref,"fasta")
alt_ref.id="REF_%s" %(str(alt_ref.id).split(".")[0])
SeqIO.write(alt_ref,"%s/%s.fna" %(Results_folder_name,alt_ref.id),"fasta")

for gen_names in genomes_query:
	ext=gen_names.strip().split(".")[-1]
	fname=".".join(gen_names.strip().split(".")[:-1])
	if (ext == "fa" or ext == "fasta") :
		os.rename ("%s/%s" %(ABS_query_folder,gen_names), "%s/%s.fna" %(ABS_query_folder,fname))

NEAREST=[]
for i in genomes_query:
	cmd="mash dist %s/%s %s -p %i 2>/dev/null | sort -gr -k5 | head -n %i | cut -f2"  %(ABS_query_folder,i,db_sketch,threads,nearest)
	NEAREST=NEAREST+(os.popen(cmd).read().strip().strip('\n').split('\n'))

NEAREST=list(set(NEAREST))
NEAREST=[i.split("/")[-1] for i in NEAREST]
print ("Sketches completed!\n")

os.chdir(Results_folder_name)
with open("query_genome_list.txt","w") as q:
	for i in genomes_query:
		q.write("%s\n" %(i))
with open("backgroud_genome_list.txt","w") as q:
	for i in NEAREST:
		q.write("%s\n" %(i))
os.mkdir("Background")

if args.bkg_folder:
	print ("Retrieving nearest genomes from the %s folder\n" %args.bkg_folder)
	os.chdir(path_dir)
	background_folder=os.path.abspath(bkg_folder)
	os.chdir(Results_folder_name)
	for N in NEAREST:
		cmd="cp %s/%s Background/" %(background_folder,N)
		os.system(cmd)
else:
	os.chdir("Background")
	print ("Retrieving nearest genomes from the BV-BRC Database...\n")
	for N in NEAREST:
		N=N.replace(".fna","").replace(".fasta","").replace(".fa","")
		cmd = "wget -qN ftp://ftp.bvbrc.org/genomes/%s/%s.fna" %(N,N)
		os.system(cmd)
	os.chdir(path_dir+"/"+Results_folder_name)

os.mkdir("Align")
os.mkdir("Analysis_Dataset")

for i in glob.glob("Background/*"):
	name=i.strip().split("/")[-1]
	oF=open("Align/BD_%s" %(name),"w")
	origsize=0
	trimmedsize=0
	for x in SeqIO.parse(i,"fasta"):
		origsize+=len(str(x.seq))
		if len(str(x.seq))>=contig_length+2*borders:
			x.seq=x.seq[borders:-borders]
			trimmedsize+=len(str(x.seq))
			SeqIO.write(x,oF,"fasta")
	oF.close()
	if trimmedsize<float(origsize)*0.90:
		print("\nDatabase genome %s was too short after contig polishing and was excluded from the Analysis Dataset (AD)" %(name))
		os.system("rm Align/BD_%s" %(name))
	else:
		os.system("cp %s Analysis_Dataset/BD_%s" %(i,name))

for i in glob.glob("%s/*" %(ABS_query_folder)):
	name=i.strip().split("/")[-1]
	oF=open("Align/%s" %(name),"w")
	origsize=0
	trimmedsize=0
	for x in SeqIO.parse(i,"fasta"):
		origsize+=len(str(x.seq))
		if len(str(x.seq))>=contig_length+2*borders:
			x.seq=x.seq[borders:-borders]
			trimmedsize+=len(str(x.seq))
			SeqIO.write(x,oF,"fasta")
	oF.close()
	if trimmedsize<float(origsize)*0.90:
		sys.exit("\n\nERROR: your query genome %s was too short after contig polishing, please check your genome quality! EXECUTION ABORTED" %(name))
	else:
		os.system("cp %s Analysis_Dataset/%s" %(i,name))

aln_folder_size=os.popen("ls Align/ | wc -l").read().strip()
os.system("rm -rf Background")

if args.amrf==True:
    check_res_vir(threads,"Analysis_Dataset","%s.fna" %(alt_ref.id))

alt_ref_abs=os.path.abspath("%s.fna" %(alt_ref.id))
Mummer_snp_call(threads,alt_ref_abs,"Align")
os.mkdir("SNPs")
os.system("mv Align/*.snp SNPs/")
os.system("mv Align/Coverage_report.txt .")
snp_num=os.popen("ls SNPs/*.snp | wc -l").read().strip()
if int(snp_num)!=int(aln_folder_size):
    sys.exit("\n\nERROR! Something went wrong! Not all .snp files were produced")
os.system("rm -rf Align")

print("Performing phylogenetic reconstruction...\n")
cmd="iqtree -s SNP_alignment.core.fasta -pre SNP_alignment -m MFP+GTR+ASC -bb 1000 -nt %i" %(threads)
os.system(cmd)

os.system("$CONDA_PREFIX/bin/Rscript %s/Snpbreaker.R SNP_alignment.core.fasta %s" %(prog_dir,args.snp_threshold))

if args.amrf==True:
    emptyAMRF=os.popen("cat summary_resistance_virulence.txt | sort | uniq | wc -l").read().strip()
    if int(emptyAMRF)<=1:
        os.system("rm summary_resistance_virulence.txt")

os.system("$CONDA_PREFIX/bin/Rscript %s/annotated_tree.R SNP_alignment.treefile" %(prog_dir))

if args.meta is not None:
	os.system("$CONDA_PREFIX/bin/Rscript %s/contact_network.R %s" %(prog_dir,ABSmeta))

###ORGANIZE OUTPUT
os.mkdir("Other_output_files")
os.system("mv SNPs Other_output_files/SNPs")
os.system("mv Analysis_Dataset Other_output_files/Analysis_Dataset/")
os.system("mv SNP_alignment* Other_output_files/")

os.system("mv backgroud_genome_list.txt Other_output_files/")
os.system("mv clusters_manual_threshold_%s_snps.csv Other_output_files/" %(args.snp_threshold))
os.system("mv query_genome_list.txt Other_output_files/")
os.system("mv snp_distance_matrix.tsv Other_output_files/")
os.system("mv SNP_genomes.list Other_output_files/")
os.system("mv SNP_positions.core.list Other_output_files/")
if args.amrf==True:
    os.system("mv summary_resistance_virulence.txt Other_output_files/")

os.chdir(path_dir)
os.system("rm %s" %(alt_ref_abs))

time=float(time.time() - start_time)
hours=int(time/3600)
minutes=((time/3600)-hours)*60
print ("\n\n\n\n")
print ("####################\n")
print ("Analysis completed in %i hours and %f minutes" %(hours,minutes))
print ("Results are stored in folder %s/%s" %(path_dir,Results_folder_name))
print ("####################\n")


import configparser
import getopt
import sys

args = sys.argv[1:]
options = "hmo:f:"
long_options = ["help", "my_file", "output=", "config="]
default_ini = "sillypy.ini"
Config = configparser.ConfigParser()
Config.read(default_ini)
try:
    args, vals = getopt.getopt(args, options, long_options)
    for arg, val in args:
        if arg in ("-h", "--help"):
            print(
                "Usage: %s"
                % """
            Simple python script: takes arguments, prints out .ini file.
            -h, --help: this help
            -o, --output= : prints out value
            -f, --config=config.ini : use config.ini instead of default
            -m : returns own file name
"""
            )
            exit()
        elif arg in ("-m", "--my_file"):
            print("Displaying file_name:", sys.argv[0])
        elif arg in ("-o", "--output"):
            print(f"-o value is: {val}")
        elif arg in ("-f", "--config"):
            Config = configparser.ConfigParser()
            Config.read(val)
except getopt.error as err:
    print(str(err))
#
print(Config.sections())
for section in Config.sections():
    for option in Config[section]:
        print(option, Config.get(section, option))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Example of CliMAF session with SeaIce data (here organized 'a la EM')

"""
# S.Senesi - may 2015

from __future__ import print_function, division, unicode_literals, absolute_import

from climaf.api import *

import os

if not atCNRM or (os.getenv("USER") != "senesi" and os.getenv("USER") != "salas"):
    print("This example does work only at CNRM for some users")
    exit(0)

# Next declaration is actually built-in
dataloc(project="EM", organization="EM", url=["dummy"])

# Next declaration is now built-in in CliMAF (and also for a number of non-ambiguous variables)
calias('EM', 'sic', missing=1.e+20)

# Define default value for some dataset facets
cdef("frequency", "monthly")
cdef("project", "EM")

# Define your dataset (a number of facets take default values)
sic = ds(simulation="NG89", variable="sic", period="198310", realm="I")

# Display the basic filenames involved in the dataset (all filenames in one single string)
# CliMAF will search them at the data location which is the most specific among all declared data locations
files = sic.baseFiles()
print(files)

# Let CliMAF generate a file with the exact dataset in its disk cache (this
# select period and/or variables, aggregate files...)
my_file = cfile(sic)
print(my_file)

# Check file size and content
os.system("ls -al " + my_file)
os.system("type ncdump && ncdump -h " + my_file)

# Stereopolar plot with explicit levels. See other plot arguments in on-line help : 'help(plot)'
fig1 = plot(sic, title="Sic 198310 2", proj="NH70", colors="0.1 10 30 60 80 90 95 97 98 99", contours=1, focus="ocean")
cshow(fig1)

# Create a dictionnary in order to simplify the way to provide constant plot args :
pa = dict(proj="NH70", colors="0.1 10 30 60 80 90 95 97 98 99", contours=1, focus="ocean")

# Mask part of the SIC field
masked_sic = ccdo(sic, operator='setrtomiss,99,100')
figm = plot(masked_sic, title="mask above 0.99", **pa)
cshow(figm)

# Remap Sea Ice data to atmospheric grid. Use default option (remapbil)
tas = ds(project="EM", simulation="GSAGNS1", variable="tas", period="1975", realm="L")
sic_on_tas_grid = regrid(sic, tas)
fig2 = plot(sic_on_tas_grid, title="regridded_sic", **pa)
cshow(fig2)

# Access daily data
sicd = ds(simulation="NG89", variable="sic", period="198303-198304", realm="I", frequency='daily')
print(sicd.baseFiles())

exit(0)