#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# -*- coding: utf-8 -*-
#
# Last modified: Fri, 14 Oct 2022 02:22:56 +0900
import numpy as np
import pandas as pd
import os

from .measurePhenotypes import measurePhenotypes

from ..util.isNotebook import isnotebook

if isnotebook():
    from tqdm.notebook import tqdm
else:
    from tqdm import tqdm

def annotateLineageIdx(**kwargs):
    if len({"matFilePath", "segImgsPath", "rawImgsPath"} & set(kwargs.keys())) > 2:
        originFrame = 0
        if "originFrame" in kwargs.keys():
            originFrame = kwargs["originFrame"]
        return annotateSchnitz(
            kwargs["matFilePath"],
            kwargs["segImgsPath"],
            kwargs["rawImgsPath"],
            originFrame,
        )
    elif "tanauchi" in kwargs.keys():
        cutOff = None
        if "cutOff" in kwargs.keys():
            cutOff = kwargs["cutOff"]
        return annotateTanauchi(kwargs["tanauchi"], cutOff)
    elif "wang" in kwargs.keys():
        cutOff = None
        if "cutOff" in kwargs.keys():
            cutOff = kwargs["cutOff"]
        return annotateWang(kwargs["wang"], cutOff)
    elif "wakamoto" in kwargs.keys():
        cutOff = None
        if "cutOff" in kwargs.keys():
            cutOff = kwargs["cutOff"]

        return annotateHashimoto(kwargs["wakamoto"], cutOff)

    else:
        print("Error")
        sys.exit(-1)

def annotateWang(WangDir, cutOff=None, fileTag=".dat"):
    coord = sorted(
        [
            os.path.join(WangDir, d)
            for d in os.listdir(WangDir)
            if os.path.isdir(os.path.join(WangDir, d)) and "xy" in d
        ]
    )

    linIdx = 0
    ID = []
    motherID = []
    uID = 0
    Z = []
    intensity = []
    area = []
    daughter1ID = []
    daughter2ID = []
    cenX = []
    cenY = []
    linIdx = []

    lineages = []
    if cutOff is not None:
        x, y = cutOff
        for dName in coord:
            xy = int(dName[-2:])
            if xy <= x:
                lins = sorted(
                    [
                        os.path.join(dName, f)
                        for f in os.listdir(dName)
                        if os.path.isfile(os.path.join(dName, f)) and fileTag in f
                    ]
                )
                lineages += lins[:y]
    else:
        for dName in coord:
            lineages += [
                os.path.join(dName, f)
                for f in os.listdir(dName)
                if os.path.isfile(os.path.join(dName, f)) and fileTag in f
            ]

    for lin in tqdm(lineages):
        cellNo = 0
        with open(lin, "r") as data:
            next(data)
            for line in data:
                if cellNo == 0:
                    motheruID = -1
                else:
                    daughter1ID.append(uID)
                    motheruID = uID - 1
                motherID.append(motheruID)

                ID.append(cellNo)
                aa = line.split(" ")

                cenX.append(float(aa[6]))
                cenY.append(float(aa[7]))
                Z.append(int(aa[0]))
                intensity.append(float(aa[5]))
                area.append(float(aa[4]))  # Acutally Length
                if int(aa[1]) == 1:
                    daughter2ID.append(-3)
                else:
                    daughter2ID.append(-1)
                cellNo += 1
                uID += 1
                linIdx.append(lineages.index(lin))
            daughter1ID.append(-1)
    cellDict = {
        "ID": np.array(ID),
        "uID": np.array(range(uID)),
        "motherID": np.array(motherID),
        "daughter1ID": np.array(daughter1ID),
        "daughter2ID": np.array(daughter2ID),
        "cenX": np.array(cenX),
        "cenY": np.array(cenY),
        "Z": np.array(Z),
        "cellNo": np.array(ID),
        "intensity": np.array(intensity),
        "area": np.array(area),
        "linIdx": np.array(linIdx),
    }
    # for key in cellDict.keys():
    #     print(key,cellDict[key])
    CellDfWP = pd.DataFrame(cellDict)

    return CellDfWP

def annotateTanauchi(TanauchiDir, cutOff=None):
    files = os.listdir(TanauchiDir)
    linIdx = 0
    ID = []
    motherID = []
    uID = 0
    Z = []
    intensity = []
    area = []
    daughter1ID = []
    daughter2ID = []
    cenX = []
    cenY = []
    linIdx = []

    lineages = []
    if cutOff is not None:
        for lin in files:
            x, y = cutOff
            coord = lin[2:].split(".")[0].split("_")
            if int(coord[0]) - 1 < x and int(coord[1]) - 1 < y:
                lineages.append(lin)
    else:
        lineages = files

    for lin in tqdm(lineages):
        cellNo = 0
        coord = lin[2:].split(".")[0].split("_")
        with open(os.path.join(TanauchiDir, lin), "r") as data:
            for line in data:
                cenX.append(int(coord[0]))
                cenY.append(int(coord[1]))
                if cellNo == 0:
                    motheruID = -1
                else:
                    daughter1ID.append(uID)
                    motheruID = uID - 1
                motherID.append(motheruID)

                ID.append(cellNo)
                aa = line.split(",")
                Z.append(int(aa[0]) - 1)
                intensity.append(float(aa[4]))
                area.append(float(aa[2]))  # Acutally Length
                if int(aa[1]) == 1:
                    daughter2ID.append(-3)
                else:
                    daughter2ID.append(-1)
                cellNo += 1
                uID += 1
                linIdx.append(lineages.index(lin))
            daughter1ID.append(-1)
    cellDict = {
        "ID": np.array(ID),
        "uID": np.array(range(uID)),
        "motherID": np.array(motherID),
        "daughter1ID": np.array(daughter1ID),
        "daughter2ID": np.array(daughter2ID),
        "cenX": np.array(cenX),
        "cenY": np.array(cenY),
        "Z": np.array(Z),
        "cellNo": np.array(ID),
        "intensity": np.array(intensity),
        "area": np.array(area),
        "linIdx": np.array(linIdx),
    }
    # for key in cellDict.keys():
    #     print(key,cellDict[key])
    CellDfWP = pd.DataFrame(cellDict)

    return CellDfWP

def annotateHashimoto(HashimotoDir, cutOff=None):
    linIdx = 0
    ID = []
    motherID = []
    uID = 0
    Z = []
    intensity = []
    area = []
    daughter1ID = []
    daughter2ID = []
    cenX = []
    cenY = []

    lineages = []

    columnName = [
        "ID",
        "uID",
        "motherID",
        "daughter1ID",
        "daughter2ID",
        "cenX",
        "cenY",
        "Z",
        "cellNo",
        "intensity",
        "area",
        "linIdx",
    ]
    CellDfWP = pd.DataFrame(columns=columnName)
    data = pd.read_table(HashimotoDir, index_col=0)
    lastCell = data[data["LastIndex"] == 1].sort_index()
    if cutOff is not None:
        lastCell = lastCell.iloc[: int(cutOff)]
    for lCell in tqdm(lastCell.iterrows()):
        progenyId, lCell = lCell
        daughter1ID = -1

        motherId = int(lCell["PreviousCell"])
        while motherId >= 0:
            info = [
                progenyId,
                motherId,
                daughter1ID,
                -1,
                float(lCell["XM"]),
                float(lCell["YM"]),
                lCell["Slice"] - 1,
                float(lCell["Mean"]) - float(lCell["Background"]),
                float(lCell["Area"]),
                linIdx,
            ]  #  Add row with this info
            columns = columnName[1:8] + columnName[-3:]  # skip over ID and cellNo
            df = dict(zip(columns, info))

            if motherId in list(CellDfWP["uID"]) and motherId != 0:
                mask = CellDfWP["uID"] == motherId
                CellDfWP.loc[mask, "daughter2ID"] = progenyId

                d1 = CellDfWP.loc[mask, "daughter1ID"]
                d2 = CellDfWP.loc[mask, "daughter2ID"]
                if int(d1) > int(d2):
                    tmp = int(d1)
                    d1 = int(d2)
                    d2 = tmp

                maskLin = CellDfWP["linIdx"] == linIdx
                linIdx -= 1
                CellDfWP.loc[maskLin, "linIdx"] = int(CellDfWP.loc[mask, "linIdx"])
                CellDfWP = CellDfWP.append(df, ignore_index=True)
                break
            #                CellDfWP[CellDfWP['uID'] == motherId]['daughter2ID'] = progenyId
            daughter1ID = progenyId
            progenyId = motherId

            if motherId == 0:
                motherId = -1
                df["motherID"] = motherId
            else:
                lCell = data.iloc[motherId - 1]
                motherId = int(lCell["PreviousCell"])

            CellDfWP = CellDfWP.append(df, ignore_index=True)
        linIdx += 1

    CellDfWP = CellDfWP.sort_values(by="Z", ascending=True)
    CellDfWP["ID"] = list(range(len(CellDfWP)))
    CellDfWP["cellNo"] = list(range(len(CellDfWP)))
    # CellDfWP['uID'] = CellDfWP['uID'].astype(int)
    # CellDfWP['motherID'] = CellDfWP['motherID'].astype(int)
    # CellDfWP['daughter1ID'] = CellDfWP['daughter1ID'].astype(int)
    # CellDfWP['daughter2ID'] = CellDfWP['daughter2ID'].astype(int)
    # CellDfWP['Z'] = CellDfWP['Z'].astype(int)
    # CellDfWP['linIdx'] = CellDfWP['linIdx'].astype(int)
    return CellDfWP

def annotateSchnitz(matFilePath, segImgsPath, rawImgsPath, originFrame=0):
    """
    Annotate lineage indices accorting to the result of cell tracking.

    Parameters
    ----------
    matFilePath : string
                  A path to MAT files which were created by Schnitzcells
    segImgsPath : string
                  A path to directory which include segmentated images which
                  were created by Schnitzcells
    rawImgsPath : string
                  A path to direcotry which include raw images
                  which were required for Schnitzcells

    Returns
    -------
    cellDfWPL : pandas.core.frame.DataFrame
                A pandas dataframe which includes tracking result and
                phenotypes of each cell, lineage indices.
                Its name is abbreviate for cell DataFrame
                With Phenotypes, Lineage indices.
                Column indices are like below
                - ID
                - uID
                - motherID
                - daughter1ID
                - daughter2ID
                - cenX
                - cenY
                - Z
                - cellNo
                - intensity
                - area
                - linIdx
    """
    cellDfWP = measurePhenotypes(matFilePath, segImgsPath, rawImgsPath, originFrame)

    numOfLin = 0
    for uID in cellDfWP["uID"]:
        daughter1ID = cellDfWP["daughter1ID"][uID]
        daughter2ID = cellDfWP["daughter2ID"][uID]
        if daughter1ID < 0 and daughter2ID < 0:
            numOfLin += 1

    linIdx = 0
    linList = np.zeros(len(cellDfWP["uID"]))

    for uID in cellDfWP["uID"]:
        motherID = cellDfWP["motherID"][uID]

        if motherID == -1:
            linList[uID] = linIdx
            linIdx += 1
        else:
            sister1ID = cellDfWP["daughter1ID"][motherID]
            sister2ID = cellDfWP["daughter2ID"][motherID]
            if sister1ID > 0 and sister2ID > 0:  # 親が分裂していたら
                if sister1ID == uID:
                    linList[uID] = linList[motherID]
                else:
                    linList[uID] = linIdx
                    linIdx += 1
            else:  # 親が分裂していなかったら
                linList[uID] = linList[motherID]

    linIdx = 0
    for uID in cellDfWP["uID"]:
        motherID = cellDfWP["motherID"][uID]

        if motherID == -1:
            linList[uID] = linIdx
            linIdx += 1
        else:
            sister1ID = cellDfWP["daughter1ID"][motherID]
            sister2ID = cellDfWP["daughter2ID"][motherID]
            if sister1ID > 0 and sister2ID > 0:  # 親が分裂していたら
                if sister1ID == uID:
                    linList[uID] = linList[motherID]
                else:
                    linList[uID] = linIdx
                    linIdx += 1
            else:  # 親が分裂していなかったら
                linList[uID] = linList[motherID]

    linList = linList.astype(np.int64)
    cellDfWP["linIdx"] = linList

    return cellDfWP

if __name__ == "__main__":
    matFilePath = (
        "/Users/itabashi/Research/Analysis/Schnitzcells/"
        "9999-99-99/488/data/488_lin.mat"
    )
    segImgsPath = (
        "/Users/itabashi/Research/Analysis/Schnitzcells/" "9999-99-99/488/segmentation/"
    )
    rawImgsPath = (
        "/Users/itabashi/Research/Experiment/microscope/"
        "2018/08/28/ECTC_8/Pos0/forAnalysis/488FS/"
    )
    cellDfWPL = annotateLineageIdx(matFilePath, segImgsPath, rawImgsPath)
    print(cellDfWPL)


import numpy as np
import netCDF4 as nd
import matplotlib.pyplot as plt
from calc_irr import *

def get_sample_points(origin, dir , ns=10):

    gps = np.zeros((origin.shape[0],dir.shape[0]))
    gps[0] = np.tan(dir/180.*np.pi) * origin[2] + origin[0]
    gps[1,:] = origin[1]*1
    gps[2,:] = 0

    ds = np.linspace(0,1,ns)
    v = np.zeros((ds.shape[0],origin.shape[0],dir.shape[0]))

    v = np.reshape(origin,(1,origin.shape[0],1)) + (gps - np.reshape(origin,(origin.shape[0],1))) * np.reshape(ds,(ds.shape[0],1,1))

    return v

def get_boxindex(p,grid):

    dxi = np.array([grid[3],grid[4],grid[5]])
    bounds = np.array([grid[0],grid[1],grid[2]])
    index = np.zeros_like(p)

    try:
        for i in range(len(index)):
            index[i] = p[i]//grid[i+3]
            if index[i] < 0:
                index[i]+=grid[i]
            if index[i] > grid[i]-1:
                index[i] -= grid[i]

    except:

        for i in range(p.shape[0]):
            for xi in range(3):
                for k in range(p.shape[2]):

                    index[i,xi,k] = p[i,xi,k]//dxi[xi]
                    if index[i,xi,k] < 0:
                        index[i,xi,k] += bounds[xi]
                    if index[i,xi,k] > bounds[xi]-1:
                        index[i,xi,k] -= bounds[xi]

    return index.astype(int)

def get_e(p,grid,fpath="job_0.183435_36.600028/mc.flx.spc.nc"):

    edirs = nd.Dataset(fpath,'r')
    Edir = np.zeros(p.shape[1])
    Edown = np.zeros(p.shape[1])
    Eup = np.zeros(p.shape[1])

    for I in range(p.shape[1]):

        i,j,k = get_boxindex(p[:,I],grid)

        Edir[I] = edirs["Edir"][j,i,k,:]
        Edown[I] = edirs["Edown"][j,i,k,:]
        Eup[I] = edirs["Eup"][j,i,k,:]

    edirs.close()

    return Edir,Edown,Eup

def get_rad(p,grid):

    Eup = np.zeros(p.shape[1])
    Edown = np.zeros_like(Eup)
    Eu,Ed = calc_Es(UMUS,PHIS,wumu,wphi,"mc.rad.spc.nc","radiance")

    for I in range(p.shape[1]):

        i,j,k = get_boxindex(p[:,I],grid)
        Eup[I] = Eu["radiance"][j,i,k,:]
        Edown[I] = Ed["radiance"][j,i,k,:]

    return Eup,Edown

UMUS = np.loadtxt("input_params.txt",dtype=str, max_rows=1)
PHIS = np.loadtxt("input_params.txt",dtype=str,skiprows=1, max_rows=1)
wumu = np.loadtxt("numus.txt", skiprows=1, max_rows=1)
wphi = np.loadtxt("nphis.txt", skiprows=1, max_rows=1)

Nx,Ny,dx,dy = np.loadtxt("input_params.txt", skiprows = 6, max_rows=1)
Zlev = np.loadtxt("input_params.txt", skiprows = 4, max_rows=1)
Nz = 2 #Zlev.shape[0]
dz = 1
sza = np.loadtxt("input_params.txt", skiprows = 2, max_rows=1)
mu = np.cos(sza/180.*np.pi)
albedo = 0.2

grid = np.array([Nx,Ny,Nz,dx,dy,dz])

cloudx = np.array([3,4])
cloudy = np.array([0])
cloudz = np.array([0,1])

camerapos = np.array([1.5,0.01,2])
camerafov = 90.
camerapixels = 90

pixeledges = np.linspace(-camerafov/2.,camerafov/2,camerapixels+1)
pixelangles = (pixeledges[0:-1] + pixeledges[1:])/2.
pixelvalues = np.zeros(pixelangles.shape)

pixelground = get_sample_points(camerapos,pixelangles)[-1]
Edir,Edown,Eup = get_e(pixelground,grid)
#Eu,Ed = get_rad(pixelground,grid)
#
#pixelvalues = Edir * albedo / np.pi + Ed*albedo/np.pi
#palt = Edir *albedo/np.pi + Edown*albedo/np.pi
#
#truth = nd.Dataset("job_panorama/mc.rad.spc.nc" , "r")

Eu,Ed = calc_Es(UMUS,PHIS,wumu,wphi,"mc.rad.spc.nc","radiance")

#fig,ax = plt.subplots(3,1)

#ax[0].plot(np.arange(pixelvalues.shape[0]),pixelvalues,label="with Edown from radiances")
#ax[1].plot(np.arange(palt.shape[0]),palt,label="with Edown from flx file")
#ax[2].plot(np.arange(truth["radiance"][0,:,0,0].shape[0]),truth["radiance"][0,:,0,0], label="from mystic panorama")
#ax[0].legend()
#ax[1].legend()
#ax[2].legend()
#
#plt.tight_layout()
#plt.show()


import json
import urllib
import requests
import types

class MovesAPIError(Exception):
    """Raised if the Moves API returns an error."""
    pass

class MovesAPINotModifed(Exception):
    """Raised if the document requested is unmodified. Need the use of etag header"""
    pass

class MovesClient(object):
    """OAuth client for the Moves API"""
    api_url = "https://api.moves-app.com/api/1.1"
    app_auth_url = "moves://app/authorize"
    web_auth_uri = "https://api.moves-app.com/oauth/v1/authorize"
    token_url = "https://api.moves-app.com/oauth/v1/access_token"
    tokeninfo_url = "https://api.moves-app.com/oauth/v1/tokeninfo"
    refresh_url = "https://api.moves-app.com/oauth/v1/access_token"

    def __init__(self, client_id=None, client_secret=None,
                 access_token=None, use_app=False):

        self.client_id = client_id
        self.client_secret = client_secret
        self.access_token = access_token
        self.auth_url = self.app_auth_url if use_app else self.web_auth_uri
        self.use_app = use_app
        self._last_headers = None

    def parse_response(self, response):
        """Parse JSON API responses."""

        return json.loads(response.text)

    def build_oauth_url(self, redirect_uri=None, scope="activity location"):
        params = {
            'client_id': self.client_id,
            'scope': scope
        }

        if not self.use_app:
            params['response_type'] = 'code'

        if redirect_uri:
            params['redirect_uri'] = redirect_uri

        # Moves hates +s for spaces, so use %20 instead.
        encoded = urllib.urlencode(params).replace('+', '%20')
        return "%s?%s" % (self.auth_url, encoded)

    def get_oauth_token(self, code, **kwargs):

        params = {
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'code': code,
            'grant_type': kwargs.get('grant_type', 'authorization_code')
        }

        if 'redirect_uri' in kwargs:
            params['redirect_uri'] = kwargs['redirect_uri']
        response = requests.post(self.token_url, params=params)
        response = json.loads(response.content)
        try:
            return response['access_token'], response['refresh_token']
        except:
            error = "<%(error)s>: %(error_description)s" % response
            raise MovesAPIError(error)

    def refresh_oauth_token(self, refresh_token, **kwargs):

        params = {
            'client_id': self.client_id,
            'client_secret': self.client_secret,
            'refresh_token': refresh_token,
            'grant_type': kwargs.get('grant_type', 'refresh_token')
        }

        response = requests.post(self.refresh_url, params=params)
        response = json.loads(response.content)
        try:
            return response['access_token'], response['refresh_token']
        except:
            error = "<%(error)s>: %(error_description)s" % response
            raise MovesAPIError(error)

    def tokeninfo(self):

        params = {
            'access_token': self.access_token
        }

        response = requests.get(self.tokeninfo_url, params=params)
        response = json.loads(response.content)
        try:
            return response
        except:
            error = "<%(error)s>: %(error_description)s" % response
            raise MovesAPIError(error)

    def api(self, path, method='GET', **kwargs):

        params = kwargs['params'] if 'params' in kwargs else {}
        data = kwargs['data'] if 'data' in kwargs else {}

        if not self.access_token and 'access_token' not in params:
            raise MovesAPIError("You must provide a valid access token.")

        url = "%s/%s" % (self.api_url, path)
        if 'access_token' in params:
            access_token = params['access_token']
            del(params['access_token'])
        else:
            access_token = self.access_token

        headers = {
            "Authorization": 'Bearer ' + access_token
        }

        if 'etag' in params:
            headers['If-None-Match'] = params['etag']
            del(params['etag'])

        resp = requests.request(method, url,
                                data=data,
                                params=params,
                                headers=headers)
        if str(resp.status_code)[0] not in ('2', '3'):
            raise MovesAPIError("Error returned via the API with status code (%s):" %
                                resp.status_code, resp.text)
        if resp.status_code == 304:
            raise MovesAPINotModifed("Unmodified")

        self._last_headers = resp.headers
        return resp

    def get(self, path, **params):
        return self.parse_response(
            self.api(path, 'GET', params=params))

    def post(self, path, **data):
        return self.parse_response(
            self.api(path, 'POST', data=data))

    def set_first_date(self):
        if not self.first_date:
            response = self.user_profile()
            self.first_date = response['profile']['firstDate']

    def __getattr__(self, name):
        '''\
Turns method calls such as "moves.foo_bar(...)" into
a call to "moves.api('/foo/bar', 'GET', params={...})"
and then parses the response.
'''
        base_path = name.replace('_', '/')

        # Define a function that does what we want.
        def closure(*path, **params):
            'Accesses the /%s API endpoints.'
            path = list(path)
            path.insert(0, base_path)
            return self.parse_response(
                self.api('/'.join(path), 'GET', params=params)
                )

        # Clone a new method with the correct name and doc string.
        retval = types.FunctionType(
            closure.func_code,
            closure.func_globals,
            name,
            closure.func_defaults,
            closure.func_closure)
        retval.func_doc =  closure.func_doc % base_path

        # Cache it to avoid additional calls to __getattr__.
        setattr(self, name, retval)
        return retval

# Give Access to last attribute
_move_client_status = ['etag', 'x-ratelimit-hourlimit', 'x-ratelimit-hourremaining',
                       'x-ratelimit-minutelimit', 'x-ratelimit-minuteremaining']
for att in _move_client_status:
    att = att.replace('-', '_')
    setattr(MovesClient, att, property(lambda self,att=att: self._last_headers.get(att, None)
                                       if self._last_headers else att))


import os, glob
from fpdf import FPDF

class Pdf_Tool:

    def __init__(self, format):
        self.pdf = FPDF(format=format)

    def save(self, dir, pdf_name):
        if not os.path.exists(dir):
            os.makedirs(dir)

        self.pdf.output(os.path.join(dir, pdf_name), "F")

    def create(self, img_path_list, dimen):
        for img_path in img_path_list:
            self.pdf.add_page()
            self.pdf.image(img_path, dimen[0], dimen[1], dimen[2], dimen[3])

if __name__ == "__main__":

    root = os.path.join(os.getcwd(), "output")
    root = "F:/E-Book/Temp/JPG"
    for i in range(1, 4):
        no = str(i).zfill(1)
        filepath_list = sorted(glob.glob(os.path.join(root, no + "/*.jpg")), key=os.path.basename)
        pdf = Pdf_Tool((2040, 1512))
        pdf.create(filepath_list, (0, 0, 2040, 1512))
        pdf.save(os.path.join(root, no), no + ".pdf")


import os
import sys

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from matplotlib import pyplot as plt
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader

# progress bar
from tqdm.auto import tqdm

# hyper parameter
train_batch_size = 64
test_batch_size = 1000
epochs = 15
lr = 0.0003
gamma = 0.7
myseed = 1220

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = torch.flatten(x, end_dim=0)
        return output

def run():
    # CUDA, macOS GPU or CPU ?
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print('Using cuda')
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print('Using mps')
    else:
        device = torch.device("cpu")
        print('Using cpu')

    # add random seed
    seed = myseed
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    # download data
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset1 = datasets.MNIST('./data', train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST('./data', train=False, transform=transform)
    train_loader = DataLoader(dataset1, batch_size=train_batch_size, shuffle=True, num_workers=1, pin_memory=True)
    test_loader = DataLoader(dataset2, batch_size=test_batch_size, shuffle=True, num_workers=1, pin_memory=True)

    model = Net().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()
    scheduler = StepLR(optimizer, step_size=1, gamma=gamma)

    # record the performance
    epoches = []
    training_accuracies = []
    training_loss = []
    testing_accuracies = []
    testing_loss = []
    training_accuracies.append('Training accuracies')
    training_loss.append('Training loss')
    testing_accuracies.append('Testing accuracies')
    testing_loss.append('Testing loss')

    # ---------- train and validate ----------
    for epoch in range(1, epochs + 1):
        total_train_loss, total_train_acc = 0.0, 0.0
        total_test_loss, total_test_acc = 0.0, 0.0
        train_batch_idx, test_batch_idx = 0, 0
        tqdm.write("[ epoch " + str(epoch) + " ]")
        epoches.append(epoch)

        # ---------- train ----------
        model.train()

        for batch in tqdm(train_loader, file=sys.stdout):

            train_batch_idx += 1
            if train_batch_idx == 1:
                tqdm.write("Training")

            imgs, labels = batch

            outputs = model(imgs.to(device))

            loss = criterion(outputs, labels.to(device))

            optimizer.zero_grad()

            loss.backward()

            nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)

            optimizer.step()

            acc = (outputs.argmax(dim=-1) == labels.to(device)).float().mean()

            total_train_loss += loss.item()
            total_train_acc += acc.item()

        train_loss = total_train_loss / train_batch_idx
        train_acc = total_train_acc / train_batch_idx

        # ---------- validate ----------
        model.eval()

        for batch in tqdm(test_loader, file=sys.stdout):
            if test_batch_idx == 1:
                tqdm.write("Testing")

            test_batch_idx += 1

            imgs, labels = batch

            with torch.no_grad():
                outputs = model(imgs.to(device))

            loss = criterion(outputs, labels.to(device))

            acc = (outputs.argmax(dim=-1) == labels.to(device)).float().mean()

            total_test_loss += loss.item()
            total_test_acc += acc.item()

        test_loss = total_test_loss / test_batch_idx
        test_acc = total_test_acc / test_batch_idx

        training_accuracies.append(train_acc)
        training_loss.append(train_loss)
        testing_accuracies.append(test_acc)
        testing_loss.append(test_loss)
        print('Training Loss:', training_loss[epoch], 'Training Accuracy:', (100 * training_accuracies[epoch]),
              '%')
        print('Testing Loss:', testing_loss[epoch], 'Testing Accuracy:', 100 * testing_accuracies[epoch], '%')
        print('')

        scheduler.step()

    plot(epoches, training_loss)
    plot(epoches, training_accuracies)
    plot(epoches, testing_accuracies)
    plot(epoches, testing_loss)

def plot(epoches, performance):
    label = performance.pop(0)
    plt.title(label)
    plt.plot(epoches, performance, label=label)
    plt.xlabel('epoches')
    plt.legend()
    plt.savefig(label + '.jpg')
    plt.show()

if __name__ == '__main__':
    run()


from phrase import Phrase
import random

class Game:
    def __init__(self):
        self.missed = 0
        self.phrases = self.create_phrases()
        self.active_phrase = self.get_random_phrase()
        self.guesses = []

    def create_phrases(self):
        phrases = [Phrase("Most things that never get done never get done because they never get started"), Phrase("Discipline equals freedom"), Phrase("Do what you can with all you have wherever you are"), Phrase("Get after it"), Phrase("Up and Rock and Roll")]
        return phrases

    def get_random_phrase(self):
        return random.choice(self.phrases)

    def welcome(self):
        print("="*60+"""
*** WELCOME TO NIKOLAI'S MOTIVATIONAL PHRASE HUNTER 2020 ***
"""+"="*60, "\nRULES ===> You've got 5 tries to guess the phrase.\nPlease enter 1 letter at a time.\n")

    def get_guess(self):
        while True:
            user_guess = (input("Please enter a letter: ")).lower()
            if not user_guess.isalpha():
                print("That's not a valid selection. Please enter a letter.")
            elif len(user_guess) != 1:
                print("Please enter one letter at a time.")
            else:
                return user_guess

    def start(self):
        self.welcome()
        self.active_phrase.display(self.guesses)

        while not self.missed >= 5:
            print(f"*** Number missed: {self.missed} *** \n")
            user_guess = self.get_guess()
            self.guesses.append(user_guess)
            if self.active_phrase.check_guess(user_guess):
                print("YAY!\n")
                self.active_phrase.display(self.guesses)
                if self.active_phrase.check_complete(self.guesses):
                    print("CONGRATS! You win!\n")
                    break
            if not self.active_phrase.check_guess(user_guess):
                self.missed += 1
                print("\nBummer :(\n")

        if self.missed == 5:
            print("You lost. Please play again!\n")


from selenium import webdriver
from selenium.webdriver.common.keys import Keys
import time

driver=webdriver.Chrome(executable_path='./driver/chromedriver')
driver.get('https://web.whatsapp.com/')

input("please scan qr code and press any key to continue:")

RM=driver.find_element_by_css_selector('span[title="Assignments & CT"]')
RM.click()

testinput=driver.find_element_by_xpath("/html/body/div/div[1]/div[1]/div[4]/div[1]/footer/div[1]/div[2]/div/div[1]/div/div[2]")
time.sleep(10)
testinput.send_keys("Hello friends")
testinput.send_keys(Keys.RETURN)


import matplotlib.pyplot as plt
import numpy as np
plt.ion()
plt.rcParams['axes.labelsize'] = 18
plt.rcParams['axes.titlesize'] = 20
plt.rcParams['font.size'] = 16
plt.rcParams['lines.linewidth'] = 2.0
plt.rcParams['lines.markersize'] = 8
plt.rcParams['legend.fontsize'] = 14

class Simulator:
    eps = 1e-16

    def __init__(self, num_agents = 15, max_iterations = 1000, step_size = None, \
                convergence_tol = 0.001, x_bounds = (0,1), y_bounds = (0, 1)):

        # convergence_tol : % of dimensions of the room
        self.convergence_tol = convergence_tol

        # Dimensions of the room
        self.x_bounds = x_bounds
        self.y_bounds = y_bounds

        self.step_size = step_size
        self.num_agents = num_agents
        self.max_iterations = max_iterations

        self.iteration = 0
        self.converged_at_iteration = None
        self.mean_step = []

        self.__initialize_positions()

        self.__choose_attractors()

    def __find_third_vertex(self, first_vertex, second_vertex):

        """ Returns both possible options for the third vertex that makes an
        equilateral triangle with two given points"""

        # Midpoint:
        mid_x, mid_y = 0.5*(first_vertex[0] + second_vertex[0]), 0.5*(first_vertex[1] + second_vertex[1])

        # Squared length of side of equilateral triangle:
        D2 = (first_vertex[0] - second_vertex[0])**2 + (first_vertex[1] - second_vertex[1])**2

        y_diff = first_vertex[1] - second_vertex[1]

        if y_diff < Simulator.eps:
            # avoid division by zero
            y_diff += Simulator.eps

        # Negative Reciprocal slope of line joining first and second vertex:
        slope = -(first_vertex[0] - second_vertex[0]) / y_diff

        # Intercept of perpendicular bisector line between first and second vertex:
        intercept = mid_y - slope * mid_x

        # For the quadratic formula:
        A = 1
        B = -2 * mid_x
        C = mid_x**2 - (3/4) * D2 /(slope**2 + 1)

        Rx = np.roots([A, B, C])
        Ry = slope*Rx + intercept

        vertex_options = (Rx, Ry)

        return vertex_options

    def __find_projections(self, target_location_x, target_location_y, current_x, current_y):
        R_vect = np.array([target_location_x - current_x, target_location_y - current_y])
        Rx_vect = np.array([target_location_x - current_x, 0])
        Ry_vect = np.array([0, target_location_y - current_y])

        # Make the distance travelled a proportion of R_vect
        x_projection = self.step_size * np.dot(Rx_vect, R_vect) / (np.linalg.norm(Rx_vect) + Simulator.eps)
        y_projection = self.step_size * np.dot(Ry_vect, R_vect) / (np.linalg.norm(Ry_vect) + Simulator.eps)

        signed_projection = np.sign(R_vect) * np.array([x_projection, y_projection])
        return (signed_projection[0], signed_projection[1])

    def __initialize_positions(self):

        # Container for the whole simulation:
        self.X = np.zeros((self.num_agents, self.max_iterations + 1))
        self.Y = np.zeros((self.num_agents, self.max_iterations + 1))

        # Initialize first positions:
        self.X[:,0] = np.random.rand(self.num_agents,)
        self.Y[:,0] = np.random.rand(self.num_agents,)

    def __choose_attractors(self):
        if self.num_agents < 3:
            raise Exception('The number of agents must be at least 3')

        # Populate the options for each agent to follow, anyone but herself
        options = np.arange(self.num_agents)
        options = np.tile(options,(len(options),1))
        options = options[~np.eye(options.shape[0],dtype=bool)].reshape(options.shape[0],-1)

        # Pick two random indices to options for two people to follow
        # (scale the random number by the range and round.)

        # Actually will need to loop here for the second agent because have to make sure not
        # choosing same two people:

        # Initialize
        follows = np.zeros((self.num_agents, 2))

        # First attractor:
        follows[:, 0, np.newaxis] = np.round( (options.shape[1] - 1) * np.random.rand(self.num_agents, 1) ).astype(int)

        # Second attractor:
        for agent in range(self.num_agents):
            firstDraw = follows[agent,0]

            # Initialize:
            secondDraw = firstDraw
            while secondDraw == firstDraw:
                # Want a different random draw from the options
                secondDraw = np.round( (options.shape[1] - 1) * np.random.rand() ).astype(int)
            follows[agent,1] = secondDraw

        follows=follows.astype(int)

        self.first_attractor = options[np.arange(options.shape[0]), follows[:,0], np.newaxis]
        self.second_attractor = options[np.arange(options.shape[0]), follows[:,1], np.newaxis]

    def _update_positions(self):
        """
        This allows each agent to jump directly to the third vertex that would create an equilateral triangle
        with the agent and the agent's two targets. However, everyone is jumping at the same time so these
        triangles are not likely to be formed until later in the simulation (if ever)
        """
        if self.step_size is not None:
            if self.step_size > 1:
                raise Exception('The step size should be less than 1')

        for agent in range(self.num_agents):

            # Find the points where you want to go to complete the triangle
            first_vertex = (self.X.item((self.first_attractor.item(agent), self.iteration)), \
               self.Y.item(self.first_attractor.item(agent), self.iteration))
            second_vertex = (self.X.item((self.second_attractor.item(agent), self.iteration)), \
               self.Y.item(self.second_attractor.item(agent), self.iteration))

            options_x, options_y = self.__find_third_vertex(first_vertex, second_vertex)

            # Find the closest of the two vertices to your current position, or the one that is inside the room:
            # For now, just don't update position if both are out of bounds

            out_of_bounds = (options_x > self.x_bounds[1]) | (options_x < self.x_bounds[0]) | \
                            (options_y > self.y_bounds[1]) | (options_y < self.y_bounds[0])

            options_x = options_x[~out_of_bounds]
            options_y = options_y[~out_of_bounds]

            current_x = self.X[agent, self.iteration]
            current_y = self.Y[agent, self.iteration]

            # Update the next position
            if len(options_x) > 1:
                # Distance to first & second options:
                D1 = ( (options_x[0] - current_x)**2 + (options_y[0] - current_y)**2 )**0.5
                D2 = ( (options_x[1] - current_x)**2 + (options_y[1] - current_y)**2 )**0.5
                closest_ind = np.argmin([D1, D2])

                if self.step_size is not None:
                    x_projection, y_projection = self.__find_projections(options_x.item(closest_ind), \
                                            options_y.item(closest_ind), current_x, current_y)
                    self.X[agent, self.iteration + 1] = current_x + x_projection
                    self.Y[agent, self.iteration + 1] = current_y + y_projection
                else:
                    self.X[agent, self.iteration + 1] = options_x[closest_ind]
                    self.Y[agent, self.iteration + 1] = options_y[closest_ind]

            elif len(options_x) == 1:
                if self.step_size is not None:
                    x_projection, y_projection = self.__find_projections(options_x.item(0), \
                                            options_y.item(0), current_x, current_y)
                    self.X[agent, self.iteration + 1] = current_x + x_projection
                    self.Y[agent, self.iteration + 1] = current_y + y_projection
                else:
                    self.X[agent, self.iteration + 1] = options_x
                    self.Y[agent, self.iteration + 1] = options_y

            else: # Don't change position
                self.X[agent, self.iteration + 1] = current_x
                self.Y[agent, self.iteration + 1] = current_y

    def plot_positions(self, initialize_plot, plot_sides = False, zoom = False):

        if initialize_plot:
            # Setting the x and y data explictly for dynamic plot update only works for plot, not scatter:
            # Going to follow the first attractor with a different color
            self.ax1.plot(self.X[0, self.iteration], self.Y[0, self.iteration], 'r.')
            self.ax1.plot(self.X[self.first_attractor.item(0), self.iteration], \
                         self.Y[self.first_attractor.item(0), self.iteration],'r+')
            self.ax1.plot(self.X[self.second_attractor.item(0), self.iteration], \
                         self.Y[self.second_attractor.item(0), self.iteration],'r+')
            self.ax1.plot(self.X[1:, self.iteration], self.Y[1:, self.iteration],'b.')

            self.ax1.set_aspect('equal')
            self.ax1.set_xlim(self.x_bounds[0], self.x_bounds[1])
            self.ax1.set_ylim(self.y_bounds[0], self.y_bounds[1])
            self.ax1.set_ylabel("Y")
            self.ax1.set_xlabel("X")
            self.ax1.set_title("Position of Agents")
        else:
            # Plot the new position
            self.ax1.set_title("Iteration = {}".format(self.iteration))
            for lin_num, line in enumerate(self.ax1.lines):
                if lin_num==0:
                    line.set_xdata(self.X[0, self.iteration])
                    line.set_ydata(self.Y[0, self.iteration])
                elif lin_num==1:
                    line.set_xdata(self.X[self.first_attractor.item(0), self.iteration - 1])
                    line.set_ydata(self.Y[self.first_attractor.item(0), self.iteration - 1])
                elif lin_num==2:
                    line.set_xdata(self.X[self.second_attractor.item(0), self.iteration - 1])
                    line.set_ydata(self.Y[self.second_attractor.item(0), self.iteration - 1])
                else:
                    line.set_xdata(self.X[1:, self.iteration])
                    line.set_ydata(self.Y[1:, self.iteration])

        self.fig.canvas.draw()
        # This is crucial for viewing the plots from the command line:
        try:
            plt.pause(0.5)
        except Exception:
            pass

        if plot_sides:
            for agent in range(self.num_agents):
                # Grab the positions for the attractors of each agent & plot the triangle in green at the end
                X_triangle = np.hstack((self.X[agent, self.iteration], \
                                        self.X[self.first_attractor.item(agent), self.iteration], \
                                        self.X[self.second_attractor.item(agent), self.iteration], \
                                        self.X[agent, self.iteration]))
                Y_triangle = np.hstack((self.Y[agent, self.iteration], \
                                        self.Y[self.first_attractor.item(agent), self.iteration], \
                                        self.Y[self.second_attractor.item(agent), self.iteration], \
                                        self.Y[agent, self.iteration]))
                self.ax1.plot(X_triangle, Y_triangle, '-g')
        if zoom:
            # Zoom In on the final positions
            self.ax1.set_xlim(0.9 * min(self.X[:, self.iteration]), 1.1 * max(self.X[:, self.iteration]))
            self.ax1.set_ylim(0.9 * min(self.Y[:, self.iteration]), 1.1 * max(self.Y[:, self.iteration]))
            self.ax1.set_aspect('equal')

    def run(self, plot_trajectories = True, plot_convergence = True):

        if plot_trajectories:
            self.fig, self.ax1 = plt.subplots(nrows = 1, ncols = 1, figsize=(8, 8)) # two axes on figure
            self.plot_positions(initialize_plot = True)

        while self.iteration < self.max_iterations:

            # Check for convergence using mean step size for all agents:
            self.mean_step.append(np.mean( ( (self.X[:, self.iteration, np.newaxis] \
                                            - self.X[:, self.iteration - 1, np.newaxis] )**2 \
                                            + (self.Y[:, self.iteration, np.newaxis] \
                                            - self.Y[:, self.iteration - 1, np.newaxis] )**2 )**0.5 ) )

            # Define convergence as once the mean step size has dropped below the threshold for 100 iterations
            # Stop the simulation once converged.
            if self.iteration > 100: # Don't bother with convergence rules unless dealing with a significant simulation
                if all( ms <= self.convergence_tol for ms in self.mean_step[self.iteration - 100: self.iteration + 1] ):
                    self.converged_at_iteration = self.iteration
                    self.X = self.X[:, np.arange(self.iteration+1)]
                    self.Y = self.Y[:, np.arange(self.iteration+1)]
                    break

            self._update_positions()

            # Update
            self.iteration += 1

            if plot_trajectories:
                self.plot_positions(initialize_plot = False)

        if plot_convergence:
            # Plot the end positions of the agents, even if we weren't plotting
            # their trajectories throughout, along with the sides of the
            # triangles and the convergence graph
            plot_sides = True

            if not plot_trajectories:
                self.fig, self.ax1  = plt.subplots(nrows = 1, ncols = 1, figsize=(8, 8))
                initialize = True
            else:
                initialize = False

            #if self.step_size is not None:
            #    zoom = True
            #else:
            #    zoom = False

            self.plot_positions(initialize, plot_sides)

            self.fig2, self.ax2  = plt.subplots(nrows = 1, ncols = 1, figsize=(8, 4))
            self.ax2.plot(self.mean_step)
            self.ax2.set_ylabel("Mean Step Size")
            self.ax2.set_xlabel("Iteration")
            self.ax2.set_title("Motion of Agents")


import cartopy
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import numpy as np

def plot_geodetic(location_to_geodetic, edxml_parser):
    # get names and geodetics for plotting
    locs = np.asarray(list(location_to_geodetic.keys()))
    geodetic_coords = np.asarray(list(location_to_geodetic.values()))
    geodetic_coords = geodetic_coords[:, [1, 0]]
    # remove any locations from geodetic that was not in parser
    loc_in_parser = np.isin(locs, edxml_parser.sorted_unique_locs)
    locs = locs[loc_in_parser]
    geodetic_coords = geodetic_coords[loc_in_parser]
    # count occurences of each location in parser
    loc_counts = np.asarray([edxml_parser.loc_to_count[loc] for loc in locs])
    # set up figure and axes
    fig = plt.figure()
    ax = plt.axes(projection=ccrs.PlateCarree())
    # zoom view around points
    min_coord = np.min(geodetic_coords, axis=0) - 5
    max_coord = np.max(geodetic_coords, axis=0) + 5
    ax.set_extent([min_coord[0], max_coord[0], min_coord[1], max_coord[1]],
                  crs=ccrs.PlateCarree())
    # add imagery
    ax.add_feature(cartopy.feature.LAND)
    ax.add_feature(cartopy.feature.OCEAN)
    # plot points
    sc = plt.scatter(geodetic_coords[:, 0], geodetic_coords[:, 1], color='#00000088', marker='o',
                     s=2*loc_counts, transform=ccrs.PlateCarree())
    # create annotation
    # code modified from:
    # https://stackoverflow.com/questions/7908636/possible-to-make-labels-appear-when-hovering-over-a-point-in-matplotlib
    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20), textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # define func to update annotations
    def update_annot(ind):
        # get position from first point
        pos = sc.get_offsets()[ind["ind"][0]]
        annot.xy = pos
        # draw box with annotations from all points from event
        text = "\n".join([locs[n] for n in ind["ind"]])
        annot.set_text(text)
        annot.get_bbox_patch().set_alpha(0.4)

    # define func to handle clicking
    def on_click(event):
        vis = annot.get_visible()
        if event.inaxes == ax:
            cont, ind = sc.contains(event)
            # update annotation if point with data clicked
            if cont:
                update_annot(ind)
                annot.set_visible(True)
                fig.canvas.draw_idle()
            # hide annotation if point without data clicked
            else:
                if vis:
                    annot.set_visible(False)
                    fig.canvas.draw_idle()

    fig.canvas.mpl_connect("button_release_event", on_click)
    # display plot
    plt.show()


import torch
import torch.nn as nn
import torch.nn.functional as F

class BasicBlock(nn.Module):
    def __init__(self, in_channel, out_channel, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)

        self.shortcut = nn.Sequential(
                nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channel)
        )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class ResFeatureNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.f1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=7, stride=1, padding=3),
            nn.BatchNorm2d(32),
            nn.ReLU())
        self.res1 = BasicBlock(32, 64, stride=2)
        self.res2 = BasicBlock(64, 128, stride=2)
        self.res3 = BasicBlock(128, 128, stride=2)
        self.res4 = BasicBlock(128, 256, stride=2)
        # self.res5 = BasicBlock(512, 1024, stride=2)
        self.flatten = nn.Flatten()

    def forward(self, x):
        o = self.f1(x)
        o = self.res1(o)
        o = self.res2(o)
        o = self.res3(o)
        o = self.res4(o)
        # o = self.res5(o)
        o = self.flatten(o)
        return o

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

        self.f = nn.Sequential(
            nn.Conv2d(1, 24, kernel_size=7, stride=1, padding=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.Conv2d(24, 64, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.Conv2d(64, 96, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(96, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.Flatten(),   # flatten directly without bottleneck
        )

    def forward(self, x):
        out = self.f(x)
        return out

class MetricNet(nn.Module):
    def __init__(self, in_dim=4096, hidden_size=512):
        super(MetricNet, self).__init__()

        self.fc = nn.Sequential(
            nn.Linear(in_dim * 2, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 2)
        )

    def forward(self, x):
        return self.fc(x)

class Projection(nn.Module):
    def __init__(self, in_dim=4096, hidden_size=1024):
        super(MetricNet, self).__init__()

        self.fc = nn.Sequential(
            nn.Linear(in_dim, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size)
        )

    def forward(self, x):
        return self.fc(x)

if __name__ == "__main__":
    x = torch.randn(1, 1, 64, 64)
    # m = FeatureNet()
    m = ResFeatureNet()
    o = m(x)
    print(o.size())
    from utils import cal_parameters
    print(cal_parameters(m))


"""Fedor's Upper Envelope algorithm.

Based on the original MATLAB code by Fedor Iskhakov:
https://github.com/fediskhakov/dcegm/blob/master/model_retirement.m

"""
from typing import Callable
from typing import Dict
from typing import List
from typing import Tuple

import numpy as np
from dcegm.interpolation import linear_interpolation_with_extrapolation
from scipy.optimize import brenth as root

eps = 2.2204e-16

def upper_envelope(
    policy: np.ndarray,
    value: np.ndarray,
    exog_grid: np.ndarray,
    state_choice_vec: np.ndarray,
    params: Dict[str, float],
    compute_utility: Callable,
) -> Tuple[np.ndarray, np.ndarray]:
    """Runs the Upper Envelope algorithm and drops sub-optimal points.
    Calculates the upper envelope over the overlapping segments of the
    decision-specific value functions, which in fact are value "correspondences"
    in this case, where multiple solutions are detected. The dominated grid
    points are then eliminated from the endogenous wealth grid.
    Discrete choices introduce kinks and non-concave regions in the value
    function that lead to discontinuities in the policy function of the
    continuous (consumption) choice. In particular, the value function has a
    non-concave region where the decision-specific values of the
    alternative discrete choices (e.g. continued work or retirement) cross.
    These are referred to as "primary" kinks.
    As a result, multiple local optima for consumption emerge and the Euler
    equation has multiple solutions.
    Moreover, these "primary" kinks propagate back in time and manifest
    themselves in an accumulation of "secondary" kinks in the choice-specific
    value functions in earlier time periods, which, in turn, also produce an
    increasing number of discontinuities in the consumption functions
    in earlier periods of the life cycle.
    These discontinuities in consumption rules in period t are caused by the
    worker's anticipation of landing exactly at the kink points in the
    subsequent periods t + 1, t + 2, ..., T under the optimal consumption policy.
    Args:
        policy (np.ndarray): Array of choice-specific consumption policy
            of shape (2, n_grid_wealth).
            Position [0, :] of the arrays contain the endogenous grid over wealth M,
            and [1, :] stores the corresponding value of the (consumption) policy
            function c(M, d), for each time period and each discrete choice.
        value (np.ndarray): Array of choice-specific value function
            of shape (2, n_grid_wealth).
            Position [0, :] of the array contains the endogenous grid over wealth M,
            and [1, :] stores the corresponding value of the value function v(M, d),
            for each time period and each discrete choice.
        exog_grid (np.ndarray): 1d array of exogenous savings grid of shape
            (n_grid_wealth,).
        choice (int): The current choice.
        params (dict): Dictionary containing the model's parameters.
        compute_value (callable): Function to compute the agent's value.

    Returns:
        (tuple) Tuple containing
        - policy_refined (np.ndarray): Worker's *refined* (consumption) policy
            function of the current period, where suboptimal points have been dropped
            and the kink points along with the corresponding interpolated values of
            the policy function have been added. Shape (2, 1.1 * n_grid_wealth).
        - value_refined (np.ndarray): Worker's *refined* value function of the
            current period, where suboptimal points have been dropped and the kink
            points along with the corresponding interpolated values of the value
            function have been added. Shape (2, 1.1 * n_grid_wealth).

    """
    n_grid_wealth = len(exog_grid)
    min_wealth_grid = np.min(value[0, 1:])
    credit_constr = False

    if value[0, 1] <= min_wealth_grid:
        segments_non_mono = locate_non_concave_regions(value)
    else:
        # Non-concave region coincides with credit constraint.
        # This happens when there is a non-monotonicity in the endogenous wealth grid
        # that goes below the first point.
        # Solution: Value function to the left of the first point is analytical,
        # so we just need to add some points to the left of the first grid point.

        credit_constr = True
        expected_value_zero_wealth = value[1, 0]

        policy, value = _augment_grid(
            policy,
            value,
            state_choice_vec,
            expected_value_zero_wealth,
            min_wealth_grid,
            n_grid_wealth,
            params,
            compute_utility=compute_utility,
        )
        segments_non_mono = locate_non_concave_regions(value)

    if len(segments_non_mono) > 1:
        _value_refined, points_to_add = compute_upper_envelope(segments_non_mono)
        index_dominated_points = find_dominated_points(
            value, _value_refined, significance=10
        )

        if credit_constr:
            value_refined = np.hstack(
                [np.array([[0], [expected_value_zero_wealth]]), _value_refined]
            )
        else:
            value_refined = _value_refined

        policy_refined = refine_policy(policy, index_dominated_points, points_to_add)

    else:
        value_refined = value
        policy_refined = policy

    # Fill array with nans to fit 10% extra grid points,
    # as the true shape is unknown ex ante
    policy_refined_with_nans = np.empty((2, int(1.1 * n_grid_wealth)))
    value_refined_with_nans = np.empty((2, int(1.1 * n_grid_wealth)))
    policy_refined_with_nans[:] = np.nan
    value_refined_with_nans[:] = np.nan

    policy_refined_with_nans[:, : policy_refined.shape[1]] = policy_refined
    value_refined_with_nans[:, : value_refined.shape[1]] = value_refined

    return policy_refined_with_nans, value_refined_with_nans

def locate_non_concave_regions(
    value: np.ndarray,
) -> List[np.ndarray]:
    """Locates non-concave regions.
    Find non-monotonicity in the endogenous wealth grid where a grid point
    to the right is smaller than its preceding point. Put differently, the
    value function bends "backwards".
    Non-concave regions in the value function are reflected by non-monotonous
    regions in the underlying endogenous wealth grid.
    Multiple solutions to the Euler equation cause the standard EGM loop to
    produce a "value correspondence" rather than a value function.
    The elimination of suboptimal grid points converts this correspondence back
    to a proper function.
    Args:
        value (np.ndarray):  Array storing the choice-specific value function
            "correspondences". Shape (2, *n_endog_wealth_grid*), where
            *n_endog_wealth_grid* is of variable length depending on the number of
            kinks and non-concave regions in the value function.
            In the presence of kinks, the value function is a "correspondence"
            rather than a function due to non-concavities.
    Returns:
        (tuple) Tuple containing:
        - value_refined (np.ndarray): Array of shape (2, *n_grid_refined*)
            containing the *refined* choice-specific value functions, which means that
            suboptimal points have been removed from the endogenous wealth grid and
            the value function "correspondence". Furthermore, kink points and the
            corresponding interpolated values of the value function have been added.
        - points_to_add (np.ndarray): Array of shape (*n_kink_points*,)
            containing the kink points and corresponding interpolated values of the
            *refined* value function that have been added to ``value_refined``.
            *n_kink_points* is of variable length.
        - index_dominated_points (np.ndarray): Array of shape (*n_dominated_points*,)
            containing the indices of dominated points in the endogenous wealth grid,
            where *n_dominated_points* is of variable length.

    """
    segments_non_mono = []

    is_monotonic = value[0, 1:] > value[0, :-1]

    niter = 0
    move_right = True

    while move_right:
        index_non_monotonic = np.where(is_monotonic != is_monotonic[0])[0]

        # Check if we are beyond the starting (left-most) point
        if len(index_non_monotonic) == 0:
            if niter > 0:
                segments_non_mono += [value]
            move_right = False
            break

        else:
            index_non_monotonic = min(index_non_monotonic)  # left-most point

            part_one, part_two = _partition_grid(value, index_non_monotonic)
            segments_non_mono += [part_one]
            value = part_two

            # Move point of first non-monotonicity to the right
            is_monotonic = is_monotonic[index_non_monotonic:]

            niter += 1

    return segments_non_mono

def compute_upper_envelope(
    segments: List[np.ndarray],
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute upper envelope and refines value function correspondence.
    The upper envelope algorithm detects suboptimal points in the value function
    correspondence. Consequently, (i) the suboptimal points are removed and the
    (ii) kink points along with their corresponding interpolated values are included.
    The elimination of suboptimal grid points converts the value
    correspondence back to a proper function. Applying both (i) and (ii)
    yields the refined endogenous wealth grid and the *refined* value function.
    Args:
        segments (List[np.ndarray]): List of non-monotonous segments in the
            endogenous wealth grid, which results in non-concavities in the
            corresponding value function. The list contains n_non_monotonous
            np.ndarrays of shape (2, *len_non_monotonous*), where
            *len_non_monotonous* is of variable length denoting the length of the
            given non-monotonous segment.
    Returns:
        (tuple) Tuple containing:

        - points_upper_env_refined (np.ndarray): Array containing the *refined*
            endogenous wealth grid and the corresponding value function.
            *refined* means suboptimal points have been dropped and the kink points
            along with the corresponding interpolated values of the value function
            have beend added.
            Shape (2, *n_grid_refined*), where *n_grid_refined* is the length of
            the *refined* endogenous grid.
        - points_to_add (np.ndarray): Array containing the kink points and
            corresponding interpolated values of the value function that have been
            added to ``points_upper_env_refined``.
            Shape (2, *n_intersect_points*), where *n_intersect_points* is the number of
            intersection points between the two uppermost segments
            (i.e. ``first_segment`` and ``second_segment``).

    """
    endog_wealth_grid = np.unique(
        np.concatenate([segments[arr][0] for arr in range(len(segments))])
    )

    values_interp = np.empty((len(segments), len(endog_wealth_grid)))
    for i, segment in enumerate(segments):
        values_interp[i, :] = linear_interpolation_with_inserting_missing_values(
            x=segment[0],
            y=segment[1],
            x_new=endog_wealth_grid,
            missing_value=-np.inf,
        )
    # values_interp has in each row the corresponding values of the upper curve
    # in the overlapping seg

    max_values_interp = np.tile(values_interp.max(axis=0), (3, 1))  # need this below
    top_segments = values_interp == max_values_interp[0, :]

    grid_points_upper_env = [endog_wealth_grid[0]]
    values_upper_env = [values_interp[0, 0]]
    intersect_points_upper_env = []
    values_intersect_upper_env = []

    move_right = True

    while move_right:
        # Index of top segment, starting at first (left-most) grid point
        index_first_segment = np.where(top_segments[:, 0])[0][0]

        for i in range(1, len(endog_wealth_grid)):
            index_second_segment = np.where(top_segments[:, i] == 1)[0][0]

            if index_second_segment != index_first_segment:
                first_segment = index_first_segment
                second_segment = index_second_segment
                first_grid_point = endog_wealth_grid[i - 1]
                second_grid_point = endog_wealth_grid[i]

                values_first_segment = (
                    linear_interpolation_with_inserting_missing_values(
                        x=segments[first_segment][0],
                        y=segments[first_segment][1],
                        x_new=np.array([first_grid_point, second_grid_point]),
                        missing_value=np.nan,
                    )
                )
                values_second_segment = (
                    linear_interpolation_with_inserting_missing_values(
                        x=segments[second_segment][0],
                        y=segments[second_segment][1],
                        x_new=np.array([first_grid_point, second_grid_point]),
                        missing_value=np.nan,
                    )
                )

                if np.all(
                    np.isfinite(
                        np.vstack([values_first_segment, values_second_segment])
                    )
                ) and np.all(np.abs(values_first_segment - values_second_segment) > 0):
                    intersect_point = root(
                        _subtract_values,
                        first_grid_point,
                        second_grid_point,
                        args=(
                            segments[first_segment],
                            segments[second_segment],
                        ),
                    )
                    value_intersect = (
                        linear_interpolation_with_inserting_missing_values(
                            x=segments[first_segment][0],
                            y=segments[first_segment][1],
                            x_new=np.array([intersect_point]),
                            missing_value=np.nan,
                        )[0]
                    )

                    values_all_segments = np.empty((len(segments), 1))
                    for segment in range(len(segments)):
                        values_all_segments[
                            segment
                        ] = linear_interpolation_with_inserting_missing_values(
                            x=segments[segment][0],
                            y=segments[segment][1],
                            x_new=np.array([intersect_point]),
                            missing_value=-np.inf,
                        )[
                            0
                        ]

                    index_max_value_intersect = np.where(
                        values_all_segments == values_all_segments.max(axis=0)
                    )[0][0]

                    if (index_max_value_intersect == first_segment) | (
                        index_max_value_intersect == second_segment
                    ):
                        # There are no other functions above
                        grid_points_upper_env.append(intersect_point)
                        values_upper_env.append(value_intersect)

                        intersect_points_upper_env.append(intersect_point)
                        values_intersect_upper_env.append(value_intersect)

                        if second_segment == index_second_segment:
                            move_right = False

            # Add point if it lies currently on the highest segment
            if (
                any(abs(segments[index_second_segment][0] - endog_wealth_grid[i]) < eps)
                is True
            ):
                grid_points_upper_env.append(endog_wealth_grid[i])
                values_upper_env.append(max_values_interp[0, i])

            index_first_segment = index_second_segment

        points_upper_env_refined = np.empty((2, len(grid_points_upper_env)))
        points_upper_env_refined[0, :] = grid_points_upper_env
        points_upper_env_refined[1, :] = values_upper_env

        points_to_add = np.empty((2, len(intersect_points_upper_env)))
        points_to_add[0] = intersect_points_upper_env
        points_to_add[1] = values_intersect_upper_env

    return points_upper_env_refined, points_to_add

def find_dominated_points(
    value_correspondence: np.ndarray,
    value_refined: np.ndarray,
    significance: int = 10,
) -> np.ndarray:
    """Returns indexes of dominated points in the value function correspondence.

    Equality is measured up to 10**(-``significance``).
    Args:
        value_correspondence (np.ndarray):  Array storing the choice-specific
            value function correspondences. Shape (2, n_endog_wealth_grid), where
            n_endog_wealth_grid is of variable length depending on the number of
            kinks and non-concave regions in the value function.
            In the presence of kinks, the value function is a correspondence
            rather than a function due to non-concavities.
        value_refined (np.ndarray): Array of refined value function, where
            suboptimal points have been dropped and kink points along with the
            corresponding interpolated values of the value function have been added.
            Shape (2, n_grid_refined), where n_grid_refined is the length of
            the refined endogenous grid.
        significance (float): Level of significance. Equality is measured up to
            10**(-``significance``).

    Returns:
        index_dominated_points (np.ndarray): Array of shape (n_dominated_points,)
            containing the indices of dominated points in the endogenous wealth grid,
            where n_dominated_points is of variable length.

    """
    sig_pos = 10**significance
    sig_neg = 10 ** (-significance)

    grid_all = np.round(value_correspondence[0, :] * sig_pos) * sig_neg
    grid_refined_sig = np.round(value_refined[0, :] * sig_pos) * sig_neg

    value_all = np.round(value_correspondence[1, :] * sig_pos) * sig_neg
    value_refined_sig = np.round(value_refined[1, :] * sig_pos) * sig_neg

    index_all = np.arange(len(grid_all))
    index_dominated_points = np.union1d(
        index_all[~np.isin(grid_all, grid_refined_sig)],
        index_all[~np.isin(value_all, value_refined_sig)],
    )

    return index_dominated_points

def refine_policy(
    policy: np.ndarray, index_dominated_points: np.ndarray, points_to_add: np.ndarray
) -> np.ndarray:
    """Drop suboptimal points from policy correspondence and add new optimal ones.

    Args:
        points_to_add (np.ndarray): Array of shape (*n_kink_points*,),
            containing the kink points and corresponding interpolated values of
            the refined value function, where *n_kink_points* is of variable
            length.
        index_dominated_points (np.ndarray): Array of shape (*n_dominated_points*,)
            containing the indices of dominated points in the endogenous wealth grid,
            where *n_dominated_points* is of variable length.

    Returns:
        (np.ndarray): Array of shape (2, *n_grid_refined*)
            containing the *refined* choice-specific policy function, which means that
            suboptimal points have been removed from the endogenous wealth grid and
            the policy "correspondence". Furthermore, kink points and the
            corresponding interpolated values of the policy function have been added.

    """
    # Remove suboptimal consumption points
    endog_wealth_grid = np.delete(policy[0, :], index_dominated_points)
    optimal_consumption = np.delete(policy[1, :], index_dominated_points)

    # Add new optimal consumption points
    new_points_policy_interp = []

    for new_grid_point in range(len(points_to_add[0, :])):
        all_points_to_the_left = np.where(
            policy[0, :] < points_to_add[0, new_grid_point]
        )[0]
        all_points_to_the_right = np.where(
            policy[0, :] > points_to_add[0, new_grid_point]
        )[0]

        last_point_to_the_left = max(
            all_points_to_the_left[
                ~np.isin(all_points_to_the_left, index_dominated_points)
            ]
        )

        # Find (scalar) point interpolated from the left
        interp_from_the_left = linear_interpolation_with_extrapolation(
            x=policy[0, :][last_point_to_the_left : last_point_to_the_left + 2],
            y=policy[1, :][last_point_to_the_left : last_point_to_the_left + 2],
            x_new=points_to_add[0][new_grid_point],
        )

        first_point_to_the_right = min(
            all_points_to_the_right[
                ~np.isin(all_points_to_the_right, index_dominated_points)
            ]
        )

        # Find (scalar) point interpolated from the right
        interp_from_the_right = linear_interpolation_with_extrapolation(
            x=policy[0, :][first_point_to_the_right - 1 : first_point_to_the_right + 1],
            y=policy[1, :][first_point_to_the_right - 1 : first_point_to_the_right + 1],
            x_new=points_to_add[0, new_grid_point],
        )

        new_points_policy_interp += [
            np.array(
                [
                    points_to_add[0, new_grid_point],
                    interp_from_the_left,
                    interp_from_the_right,
                ]
            )
        ]

    # Insert new points into the endogenous wealth grid and consumption policy
    for to_add in range(len(new_points_policy_interp)):
        index_insert = np.where(
            endog_wealth_grid > new_points_policy_interp[to_add][0]
        )[0][0]

        # 1) Add new points to policy TWICE to accurately describe discontinuities
        endog_wealth_grid = np.insert(
            endog_wealth_grid,
            index_insert,
            new_points_policy_interp[to_add][0],
        )
        endog_wealth_grid = np.insert(
            endog_wealth_grid,
            index_insert + 1,
            new_points_policy_interp[to_add][0] - 0.001 * 2.2204e-16,
        )

        # 2a) Add new optimal consumption point, interpolated from the left
        optimal_consumption = np.insert(
            optimal_consumption,
            index_insert,
            new_points_policy_interp[to_add][1],
        )
        # 2b) Add new optimal consumption point, interpolated from the right
        optimal_consumption = np.insert(
            optimal_consumption,
            index_insert + 1,
            new_points_policy_interp[to_add][2],
        )

    policy_refined = np.stack([endog_wealth_grid, optimal_consumption])

    # Make sure first element in endogenous wealth grid and optiomal consumption policy
    # are both 0.
    if policy_refined[0, 0] != 0.0:
        policy_refined = np.hstack([np.zeros((2, 1)), policy_refined])

    return policy_refined

def _augment_grid(
    policy: np.ndarray,
    value: np.ndarray,
    state_choice_vec: np.ndarray,
    expected_value_zero_wealth: np.ndarray,
    min_wealth_grid: float,
    n_grid_wealth: int,
    params,
    compute_utility: Callable,
) -> Tuple[np.ndarray, np.ndarray]:
    """Extends the endogenous wealth grid, value, and policy function to the left.

    Args:
        policy (np.ndarray):  Array storing the choice-specific
            policy function. Shape (2, *n_endog_wealth_grid*), where
            *n_endog_wealth_grid* is of variable length depending on the number of
            discontinuities in the policy function.
            In the presence of discontinuities, the policy function is a
            "correspondence" rather than a function due to multiple local optima.
        value (np.ndarray):  Array storing the choice-specific
            value function. Shape (2, *n_endog_wealth_grid*), where
            *n_endog_wealth_grid* is of variable length depending on the number of
            kinks and non-concave regions in the value function.
            In the presence of kinks, the value function is a "correspondence"
            rather than a function due to non-concavities.
        expected_value_zero_wealth (float): The agent's expected value given that she
            has a wealth of zero.
        min_wealth_grid (float): Minimal wealth level in the endogenous wealth grid.
        n_grid_wealth (int): Number of grid points in the exogenous wealth grid.
        params (dict): Dictionary containing the model's parameters.
        compute_value (callable): Function to compute the agent's value.

    Returns:
        policy_augmented (np.ndarray): Array containing endogenous grid and
            policy function with ancillary points added to the left.
            Shape (2, *n_grid_augmented*).
        value_augmented (np.ndarray): Array containing endogenous grid and
            value function with ancillary points added to the left.
            Shape (2, *n_grid_augmented*).

    """

    grid_points_to_add = np.linspace(min_wealth_grid, value[0, 1], n_grid_wealth // 10)[
        :-1
    ]

    utility = compute_utility(
        consumption=grid_points_to_add,
        params=params,
        **state_choice_vec,
    )
    values_to_add = utility + params["beta"] * expected_value_zero_wealth

    value_augmented = np.vstack(
        [
            np.append(grid_points_to_add, value[0, 1:]),
            np.append(values_to_add, value[1, 1:]),
        ]
    )
    policy_augmented = np.vstack(
        [
            np.append(grid_points_to_add, policy[0, 1:]),
            np.append(grid_points_to_add, policy[1, 1:]),
        ]
    )

    return policy_augmented, value_augmented

def _partition_grid(
    value_correspondence: np.ndarray, j: int
) -> Tuple[np.ndarray, np.ndarray]:
    """Splits the grid into two parts, 1,..., j and j, j+1,..., J.

    Note that the index ``j``, after which the separation occurs,
    is also included in the second partition.
    Args:
        value_correspondence (np.ndarray):  Array storing the choice-specific
            value function "correspondences". Shape (2, *n_endog_wealth_grid*), where
            *n_endog_wealth_grid* is of variable length depending on the number of
            kinks and non-concave regions in the value function.
            In the presence of kinks, the value function is a "correspondence"
            rather than a function due to non-concavities.
        j (int): Index denoting the location where the endogenous wealth grid is
            separated.

    Returns:
        part_one (np.ndarray): Array of shape (2, : ``j`` + 1) containing the first
            partition.
        part_two (np.ndarray): Array of shape (2, ``j``:) containing the second
            partition.

    """
    j = min(j, value_correspondence.shape[1])

    part_one = np.vstack(
        [
            value_correspondence[0, : j + 1],  # endogenous wealth grid
            value_correspondence[1, : j + 1],  # value function
        ]
    )

    # Include boundary points in both partitions
    part_two = np.vstack([value_correspondence[0, j:], value_correspondence[1, j:]])

    return part_one, part_two

def _subtract_values(grid_point: float, first_segment, second_segment):
    """Subtracts the interpolated values of the two uppermost segments."""
    values_first_segment = linear_interpolation_with_extrapolation(
        x=first_segment[0], y=first_segment[1], x_new=grid_point
    )
    values_second_segment = linear_interpolation_with_extrapolation(
        x=second_segment[0], y=second_segment[1], x_new=grid_point
    )

    diff_values_segments = values_first_segment - values_second_segment

    return diff_values_segments

def linear_interpolation_with_inserting_missing_values(x, y, x_new, missing_value):
    """Linear interpolation with inserting missing values.

    Args:
        x (np.ndarray): 1d array of shape (n,) containing the x-values.
        y (np.ndarray): 1d array of shape (n,) containing the y-values
            corresponding to the x-values.
        x_new (np.ndarray or float): 1d array of shape (m,) or float containing
            the new x-values at which to evaluate the interpolation function.
        missing_value (np.ndarray or float): Flat array of shape (1,) or float
            to set for values of x_new outside of the range of x.

    Returns:
        np.ndarray or float: 1d array of shape (m,) or float containing the
            new y-values corresponding to the new x-values.
            In case x_new contains values outside of the range of x, these
            values are set equal to missing_value.

    """
    interpol_res = linear_interpolation_with_extrapolation(x, y, x_new)
    where_to_miss = (x_new < x.min()) | (x_new > x.max())
    interpol_res[where_to_miss] = missing_value
    return interpol_res


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

from __future__ import unicode_literals

from django.test import TestCase
from django.template import Template, Context
from django.utils.html import mark_safe

from hooks.templatehook import hook
from hooks.templatetags.hooks_tags import template_hook_collect
from . import utils_hooks

class HookTagTest(TestCase):

    def setUp(self):
        self.hook_name = 'myhook'
        hook.unregister_all(self.hook_name)
        utils_hooks.myhook.unregister_all()

    def test_hook_tag(self):
        def func(context, *args, **kwargs):
            self.assertEqual(args, ("foobar", ))
            self.assertEqual(kwargs, {'bar': "bar", })
            self.assertEqual(context['foo'], "foo")
            return "hello"

        hook.register(self.hook_name, func)

        out = Template(
            "{% load hooks_tags %}"
            "{% hook hook_name 'foobar' bar='bar' %}"
        ).render(Context({"hook_name": self.hook_name, "foo": "foo", }))
        self.assertEqual(out, u"hello")

    def test_hook_tag_many(self):
        """
        Should join multiple responses
        """
        def func_a(*args, **kwargs):
            return "hello"

        def func_b(*args, **kwargs):
            return "goodbye"

        hook.register(self.hook_name, func_a)
        hook.register(self.hook_name, func_b)

        out = Template(
            "{% load hooks_tags %}"
            "{% hook hook_name 'foobar' %}"
        ).render(Context({"hook_name": self.hook_name, }))

        self.assertEqual(out, "hello\ngoodbye")

    def test_hook_tag_escaped(self):
        """
        Should escape responses (if they are not marked as safe)
        """
        def func(*args, **kwargs):
            return "hello"

        hook.register(self.hook_name, func)

        out = Template(
            "{% load hooks_tags %}"
            "{% hook hook_name 'foobar' %}"
        ).render(Context({"hook_name": self.hook_name, }))

        self.assertEqual(out, "<span>hello</span>")

    def test_hook_tag_mark_safe(self):
        """
        Should not escape safe strings
        """
        def func(*args, **kwargs):
            return mark_safe("hello")

        hook.register(self.hook_name, func)

        out = Template(
            "{% load hooks_tags %}"
            "{% hook hook_name 'foobar' %}"
        ).render(Context({"hook_name": self.hook_name, }))

        self.assertEqual(out, "hello")

    def test_template_hook_collect(self):
        def func(context, *args, **kwargs):
            self.assertEqual(context, "context")
            self.assertEqual(args, ("foo", ))
            self.assertEqual(kwargs, {'extra': "bar", })
            return "hello"

        utils_hooks.myhook.register(func)
        res = template_hook_collect(utils_hooks, 'myhook', "context", "foo", extra="bar")
        self.assertEqual(res, u"hello")

        res = template_hook_collect(utils_hooks, 'badhook')
        self.assertEqual(res, u"")

    def test_template_hook_collect_escaped(self):
        def func(*args, **kwargs):
            return "hello"

        utils_hooks.myhook.register(func)
        res = template_hook_collect(utils_hooks, 'myhook', "context", "foo", extra="bar")
        self.assertEqual(res, "<span>hello</span>")


from __future__ import unicode_literals
try:
    from urllib2 import Request
except ImportError:
    from urllib.request import Request
from mock import MagicMock, patch
from requests_kerberos import HTTPKerberosAuth

from grafana_dashboards.client.connection import (KerberosConnection,
                                                  BearerAuthConnection,
                                                  BasicAuthConnection,
                                                  SSLAuthConnection)

__author__ = 'Jakub Plichta <jakub.plichta@gmail.com>'

class Capture(object):
    """
    Class for use in method call verification that captures call argument that can be tested later on.
    """

    def __eq__(self, other):
        """
        Captures argument and always returns true to make verification successful.
        :return: True
        """
        self.value = other
        return True

def test_connection():
    connection = BasicAuthConnection('username', 'password', 'https://host')
    connection._opener = MagicMock()
    # noinspection PyProtectedMember
    connection._opener.open().read.return_value = '{"hello":"world"}'

    assert connection.make_request('/uri', {'it\'s': 'alive'}) == {'hello': 'world'}

    request = Request('https://host/uri',
                      '{"it\'s": "alive"}',
                      headers={
                          'Content-type': 'application/json',
                          'Accept': 'application/json',
                          'Authorization': b'Basic dXNlcm5hbWU6cGFzc3dvcmQ='
                      })
    capture = Capture()
    # noinspection PyProtectedMember
    connection._opener.open.assert_called_with(capture)
    assert request.get_full_url() == capture.value.get_full_url()
    assert request.header_items() == capture.value.header_items()
    assert request.get_method() == capture.value.get_method()
    assert request.data.encode('utf-8') == capture.value.data

def test_connection_with_token():
    connection = BearerAuthConnection('token', 'https://host')
    connection._opener = MagicMock()
    # noinspection PyProtectedMember
    connection._opener.open().read.return_value = '{"hello":"world"}'

    assert connection.make_request('/uri', {'it\'s': 'alive'}) == {'hello': 'world'}

    request = Request('https://host/uri',
                      '{"it\'s": "alive"}',
                      headers={
                          'Content-type': 'application/json',
                          'Accept': 'application/json',
                          'Authorization': 'Bearer token'
                      })
    capture = Capture()
    # noinspection PyProtectedMember
    connection._opener.open.assert_called_with(capture)
    assert request.get_full_url() == capture.value.get_full_url()
    assert request.header_items() == capture.value.header_items()
    assert request.get_method() == capture.value.get_method()
    assert request.data.encode('utf-8') == capture.value.data

@patch('requests.post')
def test_connection_with_kerberos(post):
    connection = KerberosConnection('https://host')

    post().json.return_value = {'hello': 'world'}

    assert connection.make_request('/uri', {'it\'s': 'alive'}) == {'hello': 'world'}

    capture = Capture()
    post.assert_called_with('https://host/uri', auth=capture, json={"it's": 'alive'}, verify=False)
    assert isinstance(capture.value, HTTPKerberosAuth)

@patch('requests.post')
def test_connection_with_sslauth(post):
    connection = SSLAuthConnection('https://host', ('/fake/cert'))

    post().json.return_value = {'hello': 'world'}

    assert connection.make_request('/uri', {'it\'s': 'alive'}) == {'hello': 'world'}

    post.assert_called_with('https://host/uri', json={"it's": 'alive'}, cert='/fake/cert')


# from time import sleep
import os
# from reply1024 import postreply1024
# import time
from datetime import datetime, timedelta

tday = datetime.now()+timedelta(hours = 8)
print(tday.hour)
tday.strftime("%Y-%m-%d %H:%M:%S")
print(tday)
if os.path.isdir("tmp")==0:
    os.mkdir("tmp")

with open("./tmp/test.txt","r+") as f:
    # f.write(f"{tday} 成功添加一行内容\n")
    con=f.read()
    a=60
    f.seek(0)
    f.truncate()
    f.write(f"{int(con)+a}")
# print(f"文件关闭了吗：{f.closed}")
with open("./tmp/test.txt") as f:
    con=f.read()

# print(f"文件关闭了吗：{f.closed}")
print("读取内容：",con)


import argparse
import torch
from model import Pretrain_SegmentationNet, DPRAN
import os
from data.dataloader import create_dataloader
from train import net_Pretrain, DPRAN_Train
import segmentation_models_pytorch as smp

def main():
    parser = argparse.ArgumentParser(description='DPRAN')
    parser.add_argument('--num_classes', default=1, type=int, help='Number of output classes [2]')
    parser.add_argument('--num_channels', default=1, type=int, help='Dimension of the input CEUS frames')
    parser.add_argument('--lr_pre', default=0.0002, type=float, help='Initial learning rate [0.0002]')
    parser.add_argument('--lr', default=0.0002, type=float, help='Initial learning rate [0.0002]')
    parser.add_argument('--num_epochs', default=50, type=int, help='Number of total training epochs [40]')
    parser.add_argument('--num_epochs_pre', default=50, type=int, help='Number of total training epochs [40]')
    parser.add_argument('--dataset', default='data', type=str, help='Dataset folder name')
    args = parser.parse_args()

    save_path = os.path.join('checkpoint')
    os.makedirs(save_path, exist_ok=True)

    layers = [32, 32, 64, 128]
    # data load and split
    train_loader, val_loader, test_loader = create_dataloader(dataset=args.dataset, batch_size=1, is_pretraining=True)

    # stage 1
    net = Pretrain_SegmentationNet(n_channels=args.num_channels, n_classes=args.num_classes, layers=layers)
    net.cuda()
    criterion = smp.losses.DiceLoss('binary', classes=None, log_loss=False, from_logits=True, smooth=0.0,
                                    ignore_index=None, eps=1e-07)
    # Optimizer
    optimizer = torch.optim.Adam(net.parameters(), lr=args.lr_pre)
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9862, last_epoch=-1)
    # Parameters
    epoch_start = 0
    epoch_end = args.num_epochs_pre
    print("Start net Pre-Training...")
    net = net_Pretrain(net, criterion, optimizer, scheduler, epoch_start, epoch_end, train_loader, val_loader,
                       save_path)
    # stage 2
    print("Start DPRAN Training...")
    model = DPRAN(n_channels=args.num_channels, n_classes=args.num_classes, layers=layers)
    model.encoder_ceus.load_state_dict(net.encoder.state_dict())
    model.cuda()

    # Optimizer and loss
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9862, last_epoch=-1)
    epoch_end = args.num_epochs
    train_loader.dataset.is_pretraining = False
    val_loader.dataset.is_pretraining = False
    test_loader.dataset.is_pretraining = False
    test_result, trained_model = DPRAN_Train(model, net, criterion, optimizer, scheduler,
                                             epoch_start, epoch_end, train_loader, val_loader,
                                             test_loader)
    torch.save({'test_rec': test_result, 'DpRAN': trained_model, 'Pretrain_SegmentationNet': net.state_dict()},
               os.path.join(save_path, 'DpRAN' + '.pt'))

if __name__ == "__main__":
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    SEED = 0
    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(SEED)
    main()


from typing import Dict, List, Union
import csv
import os
from unittest import result

def load_csv(
    file_path: str,
    delimiter: str = ',',
    has_header: bool = True,
    try_casting: bool = True

) -> List[Dict]:
    '''
    This function laods a csv file from the given path. It accepts both csv with headers and without them.

    Args:
        file_path: (str) the path to the given csv file.
        delimiter: (str) the string delimiter between columns of the csv.
        has_headers: (bool) flag to indicate if the file has headers. [Default True]

    Output:
    Returns a List of dictionaries representing each row. The keys of each dictionary ar the
    column name.

    Throws:
    - FileNotFoundError

    '''
    if not os.path.exists(file_path):
        print(f'The path {file_path} does not exists!')
        raise FileNotFoundError

    results = []

    with open(file_path, 'r') as f:
        csv_reader = csv.reader(f, delimiter=delimiter)

        if has_header:
            headers = next(csv_reader)

        for row in csv_reader:
            if try_casting:
                mapped_row = list(map(lambda item: cast_to_num(item), row))
            else:
                mapped_row = row

            new_row = { key : item for key, item in zip(headers, mapped_row)}
            results.append(new_row)

    return results

def cast_to_num(value: str) -> Union[str, int, float]:

    int_val = None
    float_val = None

    try:
        int_val = int(value)
    except ValueError:
        try:
            float_val = float(value)
        except ValueError:
            pass

    if int_val is not None:
        return int_val

    if float_val is not None:
        return float_val

    return value


# -*- coding: utf-8 -*-
"""
Created on Tue Feb 26 10:21:15 2019
This is the modl with Keras framework
@author: ago
"""

from __future__ import print_function
from matplotlib import pyplot as plt

import numpy as np
import pandas as pd
import seaborn as sns

from scipy import stats
from IPython.display import display, HTML

from sklearn import metrics
from sklearn.metrics import classification_report
from sklearn import preprocessing
from sklearn.model_selection import train_test_split
import keras
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten, Reshape
from keras.layers import Conv2D, MaxPooling2D
from keras.utils import np_utils
from keras.layers import LSTM

# Set some standard parameters upfront
pd.options.display.float_format = '{:.1f}'.format
sns.set() # Default seaborn look and feel
plt.style.use('ggplot')
print('keras version ', keras.__version__)
# Same labels will be reused throughout the program
LABELS = ['Frequency','V RMS','I ph.angle','V ph.angle','I RMS']
# The number of steps within one time segment
TIME_PERIODS = 80
# The steps to take from one segment to the next; if this value is equal to
# TIME_PERIODS, then there is no overlap between the segments
STEP_DISTANCE = 40

def read_data(file_path):

    columns = ['Time','value a','value b','label']
    df = pd.read_csv(file_path,
                     header=None,
                     names=columns)

    # will show up as NAN
    df.dropna(axis=0, how='any', inplace=True)

    return df

def convert_to_float(x):

    try:
        return np.float(x)
    except:
        return np.nan

def show_basic_dataframe_info(dataframe):

    # Shape and how many rows and columns
    print('Number of columns in the dataframe: %i' % (dataframe.shape[1]))
    print('Number of rows in the dataframe: %i\n' % (dataframe.shape[0]))

# Load data set containing all the data from csv
df = verticalStack

# Describe the data
show_basic_dataframe_info(df)
df.head(20)
# Define column name of the label vector
LABEL = 'label'
# Transform the labels from String to Integer via LabelEncoder
le = preprocessing.LabelEncoder()
# Add a new column to the existing DataFrame with the encoded values
df[LABEL] = le.fit_transform(df['label'].values.ravel())

RANDOM_SEED =50

N_TIME_STEPS = 200
N_FEATURES = 2
classes= 4
step = 1
segments = []
labels = []
for i in range(1, len(df) - N_TIME_STEPS, step):
    x1 = df['value a'].values[i: i + N_TIME_STEPS]
    x2 = df['value b'].values[i: i + N_TIME_STEPS]

    label = stats.mode(df['label'][i: i + N_TIME_STEPS])[0][0]
    segments.append([x1,x2])
    labels.append(label)

reshaped_segments = np.asarray(segments, dtype= np.float32).reshape(-1, N_TIME_STEPS, N_FEATURES)
labels = np.asarray(pd.get_dummies(labels), dtype = np.float32)

X_train, X_test, y_train, y_test = train_test_split(reshaped_segments, labels, test_size=0.2, random_state=RANDOM_SEED)

print('x_train shape: ', X_train.shape)
print( X_train.shape[0], 'training samples')
print('y_train shape: ', y_train.shape)

model = Sequential()
model.add(LSTM(200, input_shape=(X_train.shape[1], X_train.shape[2])))
model.add(Dropout(0.3))
#    model.add(LSTM(70))
#    model.add(Dropout(0.3))
model.add(Dense(classes))
callbacks_list = [
    keras.callbacks.ModelCheckpoint(
        filepath='best_model.{epoch:02d}-{val_loss:.2f}.h5',
        monitor='val_loss', save_best_only=True),
    keras.callbacks.EarlyStopping(monitor='acc', patience=1)
]

model.compile(loss='categorical_crossentropy',
                optimizer='adam', metrics=['accuracy'])
print(model.summary())

# Hyper-parameters
BATCH_SIZE = 1024
EPOCHS =10

# Enable validation to use ModelCheckpoint and EarlyStopping callbacks.
history = model.fit(X_train,
                      y_train,
                      batch_size=BATCH_SIZE,
                      epochs=EPOCHS,
                      callbacks=callbacks_list,
                      validation_split=0.2,
                      verbose=1)

plt.figure(figsize=(6, 4))
plt.plot(history.history['acc'], 'r', label='Accuracy of training data')
plt.plot(history.history['val_acc'], 'b', label='Accuracy of validation data')
plt.plot(history.history['loss'], 'r--', label='Loss of training data')
plt.plot(history.history['val_loss'], 'b--', label='Loss of validation data')
plt.title('Model Accuracy and Loss')
plt.ylabel('Accuracy and Loss')
plt.xlabel('Training Epoch')
plt.ylim(0)
plt.legend()
plt.show()

def show_confusion_matrix(validations, predictions):

    matrix = metrics.confusion_matrix(validations, predictions)
    plt.figure(figsize=(6, 4))
    sns.heatmap(matrix,
                cmap='coolwarm',
                linecolor='white',
                linewidths=1,
                xticklabels=LABELS,
                yticklabels=LABELS,
                annot=True,
                fmt='d')
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.show()

y_pred_test = model.predict(X_test)
# Take the class with the highest probability from the test predictions
max_y_pred_test = np.argmax(y_pred_test, axis=1)
max_y_test = np.argmax(y_test, axis=1)

show_confusion_matrix(max_y_test, max_y_pred_test)

print(classification_report(max_y_test, max_y_pred_test))