"""
==========================================================================
CGRA_test.py
==========================================================================
Translation for CGRAs with different configurations.

Author : Cheng Tan
  Date : Mar 2, 2019

"""

from pymtl3 import *
from pymtl3.stdlib.test             import TestSinkCL
from pymtl3.stdlib.test.test_srcs   import TestSrcRTL

from ...lib.opt_type                import *
from ...lib.messages                import *

from ...fu.flexible.FlexibleFuRTL   import FlexibleFuRTL
from ...fu.single.AdderRTL          import AdderRTL
from ...fu.single.MemUnitRTL        import MemUnitRTL
from ...fu.single.MulRTL            import MulRTL
from ...fu.single.SelRTL            import SelRTL
from ...fu.single.ShifterRTL        import ShifterRTL
from ...fu.single.LogicRTL          import LogicRTL
from ...fu.single.PhiRTL            import PhiRTL
from ...fu.single.CompRTL           import CompRTL
from ...fu.double.SeqMulAdderRTL    import SeqMulAdderRTL
from ...fu.single.BranchRTL         import BranchRTL
from ..CGRARTL                      import CGRARTL

from pymtl3.passes.backends.verilog import TranslationImportPass

#-------------------------------------------------------------------------
# Test harness
#-------------------------------------------------------------------------

class TestHarness( Component ):

  def construct( s, DUT, FunctionUnit, FuList, DataType, PredicateType,
                 CtrlType, width, height, ctrl_mem_size, data_mem_size,
                 src_opt, ctrl_waddr):

    s.num_tiles = width * height
    AddrType = mk_bits( clog2( ctrl_mem_size ) )

    s.src_opt     = [ TestSrcRTL( CtrlType, src_opt[i] )
                      for i in range( s.num_tiles ) ]
    s.ctrl_waddr  = [ TestSrcRTL( AddrType, ctrl_waddr[i] )
                      for i in range( s.num_tiles ) ]

    s.dut = DUT( DataType, PredicateType, CtrlType, width, height,
                 ctrl_mem_size, data_mem_size, len( src_opt[0] ),
                 FunctionUnit, FuList, )

    for i in range( s.num_tiles ):
      connect( s.src_opt[i].send,     s.dut.recv_wopt[i]  )
      connect( s.ctrl_waddr[i].send,  s.dut.recv_waddr[i] )

  def done( s ):
    done = True
    for i in range( s.num_tiles  ):
      if not s.src_opt[i].done():
        done = False
        break
    return done

  def line_trace( s ):
    return s.dut.line_trace()

def run_sim( test_harness, max_cycles=10 ):
  test_harness.elaborate()
  test_harness.dut.verilog_translate_import = True
  test_harness.dut.config_verilog_import = VerilatorImportConfigs(vl_Wno_list             =         ['UNSIGNED', 'UNOPTFLAT', 'WIDTH', 'WIDTHCONCAT', 'ALWCOMBORDER'])
  test_harness = TranslationImportPass()(test_harness)
  test_harness.apply( SimulationPass() )
  test_harness.sim_reset()

  # Run simulation
  ncycles = 0
  print()
  print( "{}:{}".format( ncycles, test_harness.line_trace() ))
  while not test_harness.done() and ncycles < max_cycles:
    test_harness.tick()
    ncycles += 1
    print( "{}:{}".format( ncycles, test_harness.line_trace() ))

  # Check timeout
  assert ncycles < max_cycles

  test_harness.tick()
  test_harness.tick()
  test_harness.tick()

import platform
import pytest

@pytest.mark.skipif('Linux' not in platform.platform(),
                    reason="requires linux (gcc)")
def test_cgra_universal():
  num_tile_inports  = 4
  num_tile_outports = 4
  num_xbar_inports  = 6
  num_xbar_outports = 8
  ctrl_mem_size     = 6
  width             = 2
  height            = 2
  RouteType         = mk_bits( clog2( num_xbar_inports + 1 ) )
  AddrType          = mk_bits( clog2( ctrl_mem_size ) )
  num_tiles         = width * height
  data_mem_size     = 8
  num_fu_in         = 4
  DUT               = CGRARTL
  FunctionUnit      = FlexibleFuRTL
  FuList            = [ SeqMulAdderRTL, MemUnitRTL ]#AdderRTL, MulRTL, LogicRTL, ShifterRTL, PhiRTL, CompRTL, BranchRTL, MemUnitRTL ]
  DataType          = mk_data( 32, 1 )
  PredicateType     = mk_predicate( 1, 1 )
#  FuList           = [ SeqMulAdderRTL, AdderRTL, MulRTL, LogicRTL, ShifterRTL, PhiRTL, CompRTL, BranchRTL, MemUnitRTL ]
#  DataType         = mk_data( 16, 1 )
  CtrlType          = mk_ctrl( num_fu_in, num_xbar_inports, num_xbar_outports )
  FuInType          = mk_bits( clog2( num_fu_in + 1 ) )
  pickRegister      = [ FuInType( x+1 ) for x in range( num_fu_in ) ]
  src_opt           = [ [ CtrlType( OPT_INC, b1( 0 ), pickRegister, [
                          RouteType(4), RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ),
                          CtrlType( OPT_INC, b1( 0 ), pickRegister, [
                          RouteType(4),RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ),
                          CtrlType( OPT_ADD, b1( 0 ), pickRegister, [
                          RouteType(4),RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ),
                          CtrlType( OPT_STR, b1( 0 ), pickRegister, [
                          RouteType(4),RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ),
                          CtrlType( OPT_ADD, b1( 0 ), pickRegister, [
                          RouteType(4),RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ),
                          CtrlType( OPT_ADD, b1( 0 ), pickRegister, [
                          RouteType(4),RouteType(3), RouteType(2), RouteType(1),
                          RouteType(5), RouteType(5), RouteType(5), RouteType(5)] ) ]
                          for _ in range( num_tiles ) ]
  ctrl_waddr   = [ [ AddrType( 0 ), AddrType( 1 ), AddrType( 2 ), AddrType( 3 ),
                     AddrType( 4 ), AddrType( 5 ) ] for _ in range( num_tiles ) ]
  th = TestHarness( DUT, FunctionUnit, FuList, DataType, PredicateType,
                    CtrlType, width, height, ctrl_mem_size, data_mem_size,
                    src_opt, ctrl_waddr )
  run_sim( th )


# Задача 1 Даны значения величины заработной платы заемщиков банка (zp) и значения их
# поведенческого кредитного скоринга (ks): zp = [35, 45, 190, 200, 40, 70, 54, 150, 120, 110],
# ks = [401, 574, 874, 919, 459, 739, 653, 902, 746, 832]. Используя математические
# операции, посчитать коэффициенты линейной регрессии, приняв за X заработную плату
# (то есть, zp - признак), а за y - значения скорингового балла (то есть, ks - целевая
# переменная). Произвести расчет как с использованием intercept, так и без

import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

x = np.array([35, 45, 190, 200, 40, 70, 54, 150, 120, 110]).reshape((-1, 1))
y = np.array([401, 574, 874, 919, 459, 739, 653, 902, 746, 832])

model = LinearRegression()

regres = model.fit(x,y)
b0 = regres.intercept_
b1 = regres.coef_

y_pred = model.predict(x)

print(y_pred)

plt.scatter(x,y)
plt.plot(x, b0 + b1*x )
plt.show
plt.pause(30)

# Задача 2 Посчитать коэффициент линейной регрессии при заработной плате (zp), используя
# градиентный спуск (без intercept).

x = np.array([35, 45, 190, 200, 40, 70, 54, 150, 120, 110]).reshape((-1, 1))
y = np.array([401, 574, 874, 919, 459, 739, 653, 902, 746, 832])
def mse_ (b1, y = y, x = x, n = len(x)):
    return np.sum((b1 * x - y)**2)/n

alfa = 0.000001
n = len(x)
b1 = 0.1
for i in range(1, 3000):
    b1 -= alfa * (2/n) * np.sum((b1 * x - y)*x)
    if i % 500 == 0 :
        print(f' Iteration = {i}, B1 = {b1}, MSE = {mse_(b1)}')


import openai
import pyttsx3
import speech_recognition as sr

# import OpenAI API key
openai.api_key = ""

# initialize the text-to-speech engine
tts_engine = pyttsx3.init()

def transcribe_audio_to_text(filename):
    recognizer = sr.Recognizer()
    with sr.AudioFile(filename) as source:
        audio = recognizer.record(source)
    try:
        return recognizer.recognize_google(audio)
    except:
        print("Unknown Error\n")

def generate_response(prompt):
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=prompt,
        max_tokens=4000,
        n=1,
        stop=None,
        temperature=0.5
    )
    return response["choices"][0]["text"]

def speak_text(text):
    tts_engine.say(text)
    tts_engine.runAndWait()


import os

import numpy

def picture2txt(paths: list, width: int) -> list:
    import numpy as np
    from PIL import Image
    if width <= 0:
        raise ValueError("寬度必須是正整數")
    if not paths:
        raise ValueError("路徑為空")
    file = []
    for im in paths:
        path = im
        im = Image.open(im)
        im = im.convert('L')
        height = int(width*im.size[1]/im.size[0]/2)
        im = im.resize((width, height))
        arr = np.array(im)
        arr = 255-arr
        chs0 = np.array([' ', '.', ',', '-', '^', '~', ':', ';',
                         '!', '+', '=', '*', '#', '$', '&', '@'])
        arr = arr/16
        arr = arr.astype(np.uint8)
        arr = chs0[arr]
        name_end = path.rfind('.')
        name_front = path.rfind('/', 0, name_end-1)
        name = path[name_front+1:name_end]
        file.append([name, arr])

    return file

if __name__ == '__main__':
    try:
        width = int(input("輸入txt寬度:"))
        import tkinter as tk
        from tkinter import filedialog
        tk.Tk().withdraw()
        paths = filedialog.askopenfilenames(initialdir=".", title="選取圖片", filetypes=[
            ("圖片", "png jpeg jpg"), ("所有檔案", '*')])
        file = picture2txt(paths, width)

        for name, arr in file:
            if not os.path.isdir("./輸出"):
                os.mkdir("./輸出")
            with open(f'./輸出/{name}.txt', 'w') as fp:
                for line in arr.tolist():
                    fp.write(''.join(line))
                    fp.write('\n')
            print(f'{name}.txt 《輸出完成》')

    except ValueError as e:
        print(e)
    finally:
        os.system("pause")


from sqlalchemy import create_engine
import pandas as pd
import psycopg2

def getEngine():
	try:
		connection = psycopg2.connect(dbname='eicu', user='', password='')
		# sorry for this, we need to set schema this way
		pd.read_sql_query("""
		    set search_path to eicu;
		    SELECT COUNT(*) FROM patient;
		    """, con=connection)
		return connection
	except:
		print("Unable to connect to the database")
		return None

def getConnection():
	try:
		connection = getEngine()
		return connection
	except:
		print("Unable to connect to the database")
		return None


from rq import Queue
from rq.job import Job
from rq import Retry
import yfinance
import json
# from worker import conn
from db_helpers import *
from scraper import get_sp_companies, get_high_short_interest, _parse_tickers, get_small_cap_tickers, get_dow_companies, get_nasdaq, read_russell_csv

q = get_redis_queue()

def queue_collection_tickers(tickers, collection):

    for tick in tickers:
        q.enqueue_call(
            func=write_ticker_to_collections,
            args=(tick, collection),
            retry=Retry(3),
            result_ttl=5000
        )

def large_cap():

    collection = "S&P 500"
    url = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
    tickers = _parse_tickers(*get_sp_companies(url))

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def medium_cap():

    collection = "S&P 400"
    url = "https://en.wikipedia.org/wiki/List_of_S%26P_400_companies"
    tickers = _parse_tickers(*get_sp_companies(url))

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def small_cap():
    collection = "S&P 600"
    tickers = get_small_cap_tickers()

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def dow_jones():
    collection = "Dow Jones"
    tickers = get_dow_companies()

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def nasdaq():
    collection = "Nasdaq 100"
    tickers = get_nasdaq()

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def high_short():

    collection = "High Short Interest"
    tickers = get_high_short_interest()

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def russell_3000():

    url = "https://www.ishares.com/us/products/239714/ishares-russell-3000-etf/1467271812596.ajax?fileType=csv&fileName=IWV_holdings&dataType=fund"

    collection = "Russell 3000"
    headers, data = read_russell_csv(url)

    tickers = [row[0] for row in data]

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

def russell_microcap():

    url = "https://www.ishares.com/us/products/239716/ishares-microcap-etf/1467271812596.ajax?fileType=csv&fileName=IWC_holdings&dataType=fund"

    collection = "Russell Microcap"
    headers, data = read_russell_csv(url)

    tickers = [row[0] for row in data]

    jobs = queue_collection_tickers(
        tickers,
        collection
    )

if __name__ == "__main__":

    print("Large Cap")
    large_cap()
    print("Medium Cap")
    medium_cap()
    print("Small Cap")
    small_cap()
    print("Dow Jones")
    dow_jones()
    print("Nasdaq")
    nasdaq()
    print("High Shorts")
    high_short()
    print("Russell 3000")
    russell_3000()
    print("Russell Microcap")
    russell_microcap()
    print("Finding Removed")
    find_removed()


from rest_framework import serializers
from .models import *
from django.utils import timezone
import pytz

class WeatherForecastSerializer(serializers.ModelSerializer):
    # ... existing code ...
    class Meta:
        model = Weather_Forecast
        fields = '__all__'

    def to_representation(self, instance):
        # india_tz = pytz.timezone('Asia/Kolkata')
        # timestamp = instance.timestamp.astimezone(india_tz)
        # representation = super().to_representation(instance)
        # representation['timestamp'] = timestamp.strftime('%Y-%m-%dT%H:%M:%S.%fZ')
        # return representation
        india_tz = pytz.timezone('Asia/Kolkata')
        timestamp = timezone.localtime(instance.timestamp, india_tz)
        representation = super().to_representation(instance)
        representation['timestamp'] = timestamp.isoformat()
        return representation


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Mar 11 12:20:00 2019

@author: shaunmathew

"""

from env import GridWorld
from helper import create_argparser
from helper import default_grid
import random
import time
import numpy as np

import sys

import matplotlib.pyplot as plt

def initialize(grid_world):
    states = list(grid_world.num_to_state)
    num_states = len(states)

    #Initializing epsilon soft policy using dirichlet prior
    pi = np.random.dirichlet((1,1,1,1), size=num_states)
    q_s_a = np.zeros_like(pi)

    #Initializing array that stores average return from last step and n for incremental updates
    returns = {(state,action):[0,0] for state in states for action in range(4)}

    return pi, q_s_a, returns

#Generates an episode
def gen_episode(grid_world, pi, MAX_EP_LENGTH=250000, start_func=None, **kwargs):

    starting_state = None

    if start_func:
        starting_state = start_func(**kwargs)

    #Restart environment for each ep
    grid_world.restart_env(pick_new_start=True, starting_state=starting_state)
    states = []
    actions = []
    rewards = []

    i = 1

    #Keep adding actions states and rewards to history
    while not grid_world.terminated and i <= MAX_EP_LENGTH:
        states.append(grid_world.state_to_num[grid_world.current_state])
        action = random.choices([0,1,2,3],weights=pi[grid_world.state_to_num[grid_world.current_state]],k=1)[0]
        actions.append(action)
        succ_state, reward = grid_world.move(action)
        rewards.append(reward)

        i+=1

    return states, actions, rewards

#Gets best action and chooses randomly when multiple best
def get_max_action(q_s_a, s):
    values = q_s_a[s]

    max_val = max(values)
    all_max_indices = [i for i in range(len(values)) if values[i] == max_val]
    max_action = random.choice(all_max_indices)

    return max_action

#Policy iteration and improvement
def gpi(grid_world, pi, q_s_a, returns, num_episodes=10000, sample_rate=10):
    ep_length_log = []
    time_log = []

    avg_ep_length_log = []
    avg_time_log = []

    for ep in range(num_episodes):

        start_time = time.time()
        states, actions, rewards = gen_episode(grid_world, pi)
        dist = 1

        #State action pairs
        state_actions = list(zip(states, actions))
        G = 0

        #To ease computation keep track of what is first visit so we don't have to look through entire history until time t
        first_visit = {}

        for i,sa in enumerate(state_actions):
            if sa not in first_visit:
                first_visit[sa] = i

        for t in range(len(states)-1,-1,-1):
            #G
            G = args.discount_factor*G + rewards[t]
            cur_state_action = (states[t],actions[t])
            #First visit
            if t == first_visit[cur_state_action]:

                #Incrementally updating average to compute return
                returns[cur_state_action][1] += 1
                avg = returns[cur_state_action][0] + (G - returns[cur_state_action][0])/returns[cur_state_action][1]
                returns[cur_state_action][0] = avg
                best_action = get_max_action(q_s_a, states[t])

                q_s_a[states[t],actions[t]] = avg

                #Soft policy
                for action in range(4):
                    if action == best_action:
                        pi[states[t]][action] = 1 - args.epsilon + args.epsilon/4
                    else:
                        pi[states[t]][action] = args.epsilon/4

        elapsed = (time.time() - start_time)/dist
        ep_length_log.append(len(states)/dist)
        time_log.append(elapsed)

        if (ep+1)%sample_rate==0:

            if args.debug:
                print("Episode: ", ep)
                print("Average Episode Length: ", sum(ep_length_log[-sample_rate:])/sample_rate)
                print("Average Time Per Ep : ", sum(time_log[-sample_rate:])/sample_rate)

            avg_ep_length_log.append(sum(ep_length_log[-sample_rate:])/sample_rate)
            avg_time_log.append(sum(time_log[-sample_rate:])/sample_rate)

    #Don't have to return
    return pi, q_s_a, ep_length_log, time_log, avg_ep_length_log, avg_time_log

#Outputs deterministic policy by taking max prob action
def output_deterministic_policy(pi, grid_world):
    int_to_direction = {
                0: "u",
                1: "d",
                2: "l",
                3: "r"
                }
    for row in range(len(grid_world.grid)):
        for col in range(len(grid_world.grid[row])):
            if grid_world.grid[row][col] == 0:
                print(int_to_direction[np.argmax(pi[grid_world.state_to_num[(row,col)]])], end=' ')
            elif grid_world.grid[row][col] == 3:
                print("_", end=' ')
            else:
                print(str(grid_world.grid[row][col]), end=' ')
        print("")

def main(arguments):
    parser = create_argparser()
    args = parser.parse_args(arguments)

    grid_world = GridWorld(default_grid, args.p1, args.p2)

    default_args = {"epsilon": 0.1, "discount_factor": 0.9}

    #For nice syntax
    for arg in default_args:
        if arg not in args:
            setattr(args,arg,default_args[arg])

    num_episodes = args.num_episodes

    run_dict = {}

    #injecting into global scope
    globals()['args'] = args

    num_runs = 3 if args.AVERAGE_RUNS else 1

    for i in range(num_runs):
        start_time = time.time()
        pi, q_s_a, returns = initialize(grid_world)
        _,_, ep_length_log, time_log, avg_ep_length_log, avg_time_log = gpi(grid_world, pi, q_s_a, returns, num_episodes=num_episodes)
        total_time = time.time() - start_time

        run_dict[i] = {"Episode Length": ep_length_log, "Time Per Episode": time_log, "Total Time": total_time, "Average Time Log": avg_time_log, "Average Ep Length": avg_ep_length_log}

        print("\nTook {}s to finish {} episodes".format(total_time, num_episodes))

    average_ep_lengths = np.average(np.array([run_dict[key]["Episode Length"] for key in run_dict]), axis=0)
    average_ep_time = np.average(np.array([run_dict[key]["Time Per Episode"] for key in run_dict]), axis=0)
    average_time = np.average(np.array([run_dict[key]["Total Time"] for key in run_dict]), axis=0)
    average_avg_time_log = np.average(np.array([run_dict[key]["Average Time Log"] for key in run_dict]), axis=0)
    average_avg_ep_length = np.average(np.array([run_dict[key]["Average Ep Length"] for key in run_dict]), axis=0)

    res = [average_ep_lengths, average_ep_time, average_time, average_avg_time_log, average_avg_ep_length]

    graph_names = ["Episode Length", "Time Per Episode", "Total Time in Seconds", "Time Per Episode (Moving Average 10 ep)", "Episode Length (Moving Average 10 ep)"]
    y_axis_names = ["Episode Length in Steps", "Time Per Episode in Seconds", "Total Time in Seconds", "Time Per Episode in Seconds", "Episode Length in Steps"]

    #outputting policy
    output_deterministic_policy(pi, grid_world)

    for i in [0,1]:
        t = np.linspace(1, num_episodes, num=num_episodes)[0::10]
        plt.plot(t, res[i][0::10], label="mc")
        plt.title(graph_names[i])
        plt.xlabel("Episode Number")
        plt.ylabel(y_axis_names[i])

        plt.legend()
        plt.savefig(graph_names[i] + "_mc" + ".jpg")
        plt.close()

    for i in [-2,-1]:
        t = np.linspace(1, num_episodes, num=num_episodes/10)
        plt.plot(t, res[i], label="mc")
        plt.title(graph_names[i])
        plt.xlabel("Episode Number")
        plt.ylabel(y_axis_names[i])

        plt.legend()
        plt.savefig(graph_names[i] + "_mc" + ".jpg")
        plt.close()

    return average_ep_lengths, average_ep_time, average_time, average_avg_time_log, average_avg_ep_length

if __name__ == "__main__":
    main(sys.argv[1:])

    #print(q_s_a)


from typing import TYPE_CHECKING, Any, Dict, List, Type, TypeVar

from attrs import define, field

if TYPE_CHECKING:
    from ..models.pipeline_descr import PipelineDescr
    from ..models.pipeline_runtime_state import PipelineRuntimeState

T = TypeVar("T", bound="Pipeline")

@define
class Pipeline:
    """State of a pipeline, including static configuration
    and runtime status.

        Attributes:
            descriptor (PipelineDescr): Pipeline descriptor.
            state (PipelineRuntimeState): Runtime state of the pipeine.
    """

    descriptor: "PipelineDescr"
    state: "PipelineRuntimeState"
    additional_properties: Dict[str, Any] = field(init=False, factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        descriptor = self.descriptor.to_dict()

        state = self.state.to_dict()

        field_dict: Dict[str, Any] = {}
        field_dict.update(self.additional_properties)
        field_dict.update(
            {
                "descriptor": descriptor,
                "state": state,
            }
        )

        return field_dict

    @classmethod
    def from_dict(cls: Type[T], src_dict: Dict[str, Any]) -> T:
        from ..models.pipeline_descr import PipelineDescr
        from ..models.pipeline_runtime_state import PipelineRuntimeState

        d = src_dict.copy()
        descriptor = PipelineDescr.from_dict(d.pop("descriptor"))

        state = PipelineRuntimeState.from_dict(d.pop("state"))

        pipeline = cls(
            descriptor=descriptor,
            state=state,
        )

        pipeline.additional_properties = d
        return pipeline

    @property
    def additional_keys(self) -> List[str]:
        return list(self.additional_properties.keys())

    def __getitem__(self, key: str) -> Any:
        return self.additional_properties[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.additional_properties[key] = value

    def __delitem__(self, key: str) -> None:
        del self.additional_properties[key]

    def __contains__(self, key: str) -> bool:
        return key in self.additional_properties


from django.contrib import admin
from django.urls import path, include
from django.contrib.auth import views as auth_views

urlpatterns = [

    path('admin/', admin.site.urls),

    path('', include('core.urls')),

    path('login/', auth_views.LoginView.as_view(), name='login'),

    path('logout/', auth_views.LogoutView.as_view(), {'next_page': '/'},
         name='logout')

]


from tkinter import *
from data import model

class Play:
    def __init__(self, window, text, exit_func):
        self.model = model.Model(type='fill_from_creation', text=text)
        self.load_creation_local(window, exit_func)

    def load_creation_local(self, window, exit_func):
        frame = Frame(window, bg="#d6d6d6",pady=20, padx=20)
        title = Label(frame, text='Fill in by number!', fg='#121212', bg='#d6d6d6', font=('Lucida Sans', 24), pady=(10))
        grid = Frame(frame)
        for x in range(20):
            for y in range(20):
                btn = Button(grid)
                btn.config(text=self.model.get_cell_value(x,y),bg=self.model.get_grid_cell_color(x,y), height=1, width=2, bd=1, command=lambda btn=btn: self.change_cell_color(btn))
                btn.grid(column=y, row=x)

        colors = Frame(frame, background="#d6d6d6")
        colors_title = Label(colors, text='Colors', fg='#121212', bg='#d6d6d6', font=('Lucida Sans', 16), pady=(5))
        colors_title.pack()
        self.rads = []
        for i in range (0, 10):
            rad = Radiobutton(colors, text=i, font=('Lucida Sans', 14), fg=self.model.getToolbarColor(i), activebackground=self.model.getToolbarColor(i), selectcolor=self.model.getToolbarColor(i), indicator=0, value=i, variable=self.model.selected_color, width=10)
            rad.pack(pady=5)
            self.rads.append(rad)

        menu_btns = Frame(frame, background="#d6d6d6")
        exit_btn = Button(menu_btns, text="Exit", fg='#121212', bg='#d6d6d6', activebackground='#d6d6d6', font=('Lucida Sans', 18), bd=0, pady=(10),command=exit_func)
        complete_btn = Button(menu_btns, text="Complete", fg='#121212', bg='#d6d6d6', activebackground='#d6d6d6', font=('Lucida Sans', 18), bd=0, pady=(10), command=lambda: self.complete(window, exit_func))
        exit_btn.grid(column=0, row= 0)
        complete_btn.grid(column=1, row=0)
        title.grid(column=0, row=0)
        grid.grid(column=0, row=1, padx=15)
        colors.grid(column=1, row=1, padx=5)
        menu_btns.grid(column=0, row=2)
        frame.pack()

    # Change color in model and on the grid
    def change_cell_color(self, btn):
        # print(self.model.colors[self.model.getSelectedColor()], self.model.get_grid_cell_color(btn.grid_info()["row"], btn.grid_info()["column"]))
        if self.model.getSelectedColor() == self.model.get_cell_value(btn.grid_info()["row"], btn.grid_info()["column"]):
            self.model.change_grid_cell(btn.grid_info()["row"], btn.grid_info()["column"])
            btn.config(bg=self.model.getToolbarColor(self.model.selected_color.get()))

    def complete(self, window, exit):
        if self.model.check_completed_play():
            self.show_completed(window,exit)

    def show_completed(self, window, exit):
        try:
            for key in window.children:
                window.children[key].destroy()
        except:
            pass
        finally:
            frame = Frame(window, bg="#d6d6d6",pady=5, padx=20)
            title = Label(frame, text='Nice job!', fg='#121212', bg='#d6d6d6', font=('Lucida Sans', 24), pady=(10))
            exit_btn = Button(frame, text="Exit", fg='#121212', bg='#d6d6d6', activebackground='#d6d6d6', font=('Lucida Sans', 18), bd=0, pady=(5),command=exit)
            grid = Frame(frame)
            for x in range(20):
                for y in range(20):
                    btn = Button(grid)
                    btn.config(bg=self.model.get_grid_cell_color(x,y), height=1, width=2, bd=0)
                    btn.grid(column=y, row=x)
            title.pack()
            grid.pack()
            exit_btn.pack()
            frame.pack()


# -*- coding: utf-8 -*-
"""
Created on Fri Nov 24 14:37:41 2017

@author: john
"""

# This is the class for the scrolling background
import pygame
import pygame.locals

class SpaceBackground:
    def __init__(self, screenheight, imagefile):
        self.shape = pygame.transform.scale((pygame.image.load(imagefile)),(480,640))
        self.coord = [0,0]
        self.coord2 = [0, -screenheight]
        self.y_original = self.coord[1]
        self.y2_original = self.coord2[1]

    def Show(self, surface):
        surface.blit(self.shape, self.coord)
        surface.blit(self.shape, self.coord2)

    def Scroll(self, speed_y, time):
        distance_y = speed_y * time
        self.coord[1] += distance_y
        self.coord2[1] += distance_y
        if self.coord2[1] >= 0:
            self.coord[1] = self.y_original
            self.coord2[1] = self.y2_original


# -*- coding: utf-8 -*-
import scrapy
from zero.items import IndexStatisticItem

def linesToDict(lines):
    """
    把逗号分隔的行，转换成字典
    """
    res={}
    for line in lines:
        if '=' in line:
            k,v=line.replace('var','').split('=')
            k=k.strip()
            v=v.replace('"','').replace('\r','')
            res[k]=v
    return res

kvmaping={
    #'pushDate':'zsgz00',
    #          静态市盈率      ,动态市盈率,    ,市净率        ,股息率         ,去年底静态市盈率  ,去年底滚动市盈率   ,去年底市净率
    '上证指数' :{'spe':'zsgz11','dpe':'zsgz12','pb':'zsgz13','dp':'zsgz18','lyspe':'zsgz14','lydpe':'zsgz15','lypb':'zsgz16'},
    '上证180' :{'spe':'zsgz21','dpe':'zsgz22','pb':'zsgz23','dp':'zsgz28','lyspe':'zsgz24','lydpe':'zsgz25','lypb':'zsgz26'},
    '上证50'  :{'spe':'zsgz31','dpe':'zsgz32','pb':'zsgz33','dp':'zsgz38','lyspe':'zsgz34','lydpe':'zsgz35','lypb':'zsgz36'},
    '沪深300' :{'spe':'zsgz41','dpe':'zsgz42','pb':'zsgz43','dp':'zsgz48','lyspe':'zsgz44','lydpe':'zsgz45','lypb':'zsgz46'},
    '深证成指' :{'spe':'zsgz51','dpe':'zsgz52','pb':'zsgz53','dp':'zsgz58','lyspe':'zsgz54','lydpe':'zsgz55','lypb':'zsgz56'},
    '深证100R':{'spe':'zsgz61','dpe':'zsgz62','pb':'zsgz63','dp':'zsgz68','lyspe':'zsgz64','lydpe':'zsgz65','lypb':'zsgz66'},
    '中小板指':{'spe':'zsgz71','dpe':'zsgz72','pb':'zsgz73','dp':'zsgz78','lyspe':'zsgz74','lydpe':'zsgz75','lypb':'zsgz76'},
    '上证380' :{'spe':'zsgz81','dpe':'zsgz82','pb':'zsgz83','dp':'zsgz88','lyspe':'zsgz84','lydpe':'zsgz85','lypb':'zsgz86'},
    '红利指数':{'spe':'zsgz91','dpe':'zsgz92','pb':'zsgz93','dp':'zsgz98','lyspe':'zsgz94','lydpe':'zsgz95','lypb':'zsgz96'},
    '中证红利':{'spe':'zsgz101','dpe':'zsgz102','pb':'zsgz103','dp':'zsgz108','lyspe':'zsgz104','lydpe':'zsgz105','lypb':'zsgz106'},
    '中证500' :{'spe':'zsgz111','dpe':'zsgz112','pb':'zsgz113','dp':'zsgz118','lyspe':'zsgz114','lydpe':'zsgz115','lypb':'zsgz116'}
}

def genItem(kvs,indexName):
    item             = IndexStatisticItem()
    item['push_date'] =kvs['zsgz00']
    item['index_name']=indexName
    variable=kvmaping[indexName]
    spe    =variable['spe']
    dpe    =variable['dpe']
    pb     =variable['pb']
    dp     =variable['dp']
    lyspe  =variable['lyspe']
    lydpe  =variable['lydpe']
    lypb   =variable['lypb']
    item['spe']=kvs[spe]
    item['dpe']=kvs[dpe]
    item['pb'] =kvs[pb]
    item['dp'] =kvs[dp]
    item['lyspe']=kvs[lyspe]
    item['lydpe']=kvs[lydpe]
    item['lypb']=kvs[lypb]
    return item

class IndexstatisticspiderSpider(scrapy.Spider):
    name = 'indexStatisticSpider'
    allowed_domains = ['www.csindex.com.cn']
    start_urls = ['http://www.csindex.com.cn/data/js/show_zsgz.js?str=z3l50GN6FTsOxMrb']

    def parse(self, response):
        dicts=linesToDict(response.text.split('\n'))
        for k in kvmaping:
            yield genItem(dicts,k)


from collections import OrderedDict, defaultdict

import numpy as np
import pytest

from energypy.registry import make
"""
- test cursor position (through info dict)
- test the shapes of obs & action spaces

"""

test_cases = (
    #  full charge for three steps
    (
        {'initial_charge': 0.0, 'power': 2.0, 'capacity': 100, 'episode_length': 3, 'first_reset': 'train'},
        [1.0, 1.0, 1.0],
        [2.0/2, 4.0/2, 6.0/2]
    ),
    #  full, half then full charge for three steps
    (
        {'initial_charge': 0.0, 'power': 2.0, 'capacity': 100, 'episode_length': 3, 'first_reset': 'train'},
        [1.0, 0.5, 1.0],
        [2.0/2, 3.0/2, 5.0/2]
    ),
    (
    #   discharge, charge, discharge
        {'initial_charge': 0.0, 'power': 2.0, 'capacity': 100, 'episode_length': 3, 'first_reset': 'train'},
        [-1.0, 1.0, -1.0],
        [0.0, 2.0/2, 0.0]
    )
)

@pytest.mark.parametrize('cfg, actions, expected_charges', test_cases)
def test_one_battery_charging(cfg, actions, expected_charges):
    env = make('battery', **cfg, n_batteries=1)
    env.reset()

    results = defaultdict(list)
    for action in actions:
        action = np.array(action).reshape(1, 1)
        next_obs, reward, done, info = env.step(action)
        results['charge'].append(info['final_charge'])

    assert done
    charges = np.squeeze(np.array(results['charge']))
    np.testing.assert_array_almost_equal(charges, expected_charges)

def test_battery_init():
    env = make(
        'battery',
        dataset={'name': 'random-dataset', 'n_features': 16}
    )
    #  can check shapes of dataset, action space etc

def test_many_battery_step():
    cfgs = defaultdict(list)

    actions, charges = [], []
    for test_case in test_cases:

        #  the config dict
        for k, v in test_case[0].items():
            cfgs[k].append(v)

        actions.append(test_case[1])
        charges.append(test_case[2])

    cfgs['episode_length'] = 3
    #  actions = (3, 3)
    #  needs to be timestep first!
    actions = np.array(actions).T
    expected_charges = np.array(charges).T

    env = make(
        'battery',
        n_batteries=len(test_cases),
        **cfgs,
        dataset={'name': 'random-dataset', 'n_features': 10}
    )

    #  test 1
    np.testing.assert_array_equal(cfgs['power'], env.power[0, 0])
    assert env.power.shape == (len(test_cases), 1)

    obs = env.reset()
    results = defaultdict(list)
    for action in actions:
        action = np.array(action).reshape(len(test_cases), 1)
        next_obs, reward, done, info = env.step(action)
        print(env.charge, 'charge')
        results['charge'].append(info['final_charge'])
        #  1 for the charge variable added onto our 10 features
        assert next_obs['features'].shape == (len(test_cases), 10+1)

    assert done.all()
    np.testing.assert_array_almost_equal(
        np.squeeze(results['charge']),
        np.squeeze(expected_charges)
    )


#!/usr/bin/python
# -*- coding:UTF-8 -*-
import curses
from random import randrange, choice
from collections import defaultdict

# define the action list
actions = ['Up', 'Left', 'Down', 'Right', 'Exit', 'Restart']
# map the key to action
key_codes = [ord(item) for item in 'wasdqrWASDQR']
action_dict = dict(zip(key_codes, actions*2))

# get user input
def get_user_action(keyboard):
    action = 'N'
    while action not in action_dict:
        action = keyboard.getch()
    return action_dict[action]

# transpose the matrix
def transpose(field):
    return [list(row) for row in zip(*field)]

#invert the matrix
def invert(field):
    return [row[::-1] for row in field]

# GameField class
class GameField(object):
    def __init__(self, height=4, width=4, win_value=2048):
        self.height = height
        self.width = width
        self.win_value = win_value
        self.score = 0
        self.highscore = 0
        self.reset()

    # produce a new number on the game field
    def spawn(self):
        new_element = 2 if randrange(100) > 89 else 4
        (i, j) = choice([(i, j) for i in range(self.height) for j in range(self.width) if self.field[i][j] == 0])
        self.field[i][j] = new_element

    # reset the game field
    def reset(self):
        if self.score > self.highscore:
            self.highscore = self.score
        self.score = 0
        self.field = [[0 for i in range(self.width)] for j in range(self.height)]
        self.spawn()
        self.spawn()

    # jurge whether we can move the gamefield
    def move_is_possible(self, direction):
        # jurge a row can be move left
        def row_is_left_movable(row):
            def change(i):
                if row[i]==0 and row[i+1]!=0:
                    return True
                elif row[i]!=0 and row[i]==row[i+1]:
                    return True
                else:
                    return False
            return any([change(i) for i in range(len(row)-1)])
        check={}
        check['Left'] = lambda field: any([row_is_left_movable(row) for row in field])
        check['Right'] = lambda field: check['Left'](invert(field))
        check['Up'] = lambda field: check['Left'](transpose(field))
        check['Down'] = lambda field: check['Right'](transpose(field))
        if direction in check:
            return check[direction](self.field)
        else:
            return False

    # move action
    def move(self, direction):
        def move_row_left(row):

            def tighten(row):
                new_row = [item for item in row if item != 0]
                new_row += [0 for i in range(len(row) - len(new_row))]
                return new_row

            def merge(row):
                flags = False
                new_row = []

                for i in range(len(row)):
                    if flags:
                        new_row.append(2*row[i])
                        self.score += 2*row[i]
                        flags = False
                    else:
                        if i+1 < len(row) and row[i] == row[i+1]:
                            flags = True
                            new_row.append(0)
                        else:
                            new_row.append(row[i])
                assert(len(new_row) == len(row))
                return new_row
            return tighten(merge(tighten(row)))
        moves = {}
        moves['Left'] = lambda field: [move_row_left(row) for row in field]
        moves['Up'] = lambda field: transpose([move_row_left(row) for row in transpose(field)])
        moves['Right'] = lambda field: invert(moves['Left'](invert(field)))
        moves['Down'] = lambda field: transpose(moves['Right'](transpose(field)))
        if direction in moves:
            if self.move_is_possible(direction):
                self.field = moves[direction](self.field)
                self.spawn()
                return True
            else:
                return False

    def is_win(self):
        return any(any(i >= self.win_value for i in row) for row in self.field)

    def is_gameover(self):
        return not any(self.move_is_possible(direction) for direction in actions)

    def draw(self, screen):
        help_str1 = 'W(Up) A(Left) S(Down) D(Right)'
        help_str2 = '     R(restart) Q(quit)'
        gameover_str = '      Game Over'
        win_str = '        you win'
        def cast(string):
            screen.addstr(string + '\n')

        def draw_seperator():
            line = '+' + ('+------' * self.width + '+')[1:]
            seperator = defaultdict(lambda :line)
            if not hasattr(draw_seperator, "counter"):
                draw_seperator.counter = 0
            cast(seperator[draw_seperator.counter])
            draw_seperator.counter += 1

        def draw_row(row):
            cast(''.join('|{: ^5} '.format(num) if num >0 else '|      ' for num in row)+'|')

        screen.clear()
        cast('Score:'+str(self.score))
        if self.highscore != 0:
            cast('HighScore:'+str(self.highscore))
        for row in self.field:
            draw_seperator()
            draw_row(row)
        draw_seperator()
        if self.is_win():
            cast(win_str)
        elif self.is_gameover():
            cast(gameover_str)
        else:
            cast(help_str1)
            cast(help_str2)

def main(stdscr):
    curses.use_default_colors()
    game_field = GameField(win_value=2048)
    def init():
        game_field.reset()
        return 'Game'

    def game():
        game_field.draw(stdscr)
        action = get_user_action(stdscr)
        if action == 'Restart':
            return 'Init'
        if action == 'Exit':
            return 'Exit'
        if game_field.move(action):
            if game_field.is_win():
                return 'Win'
            if game_field.is_gameover():
                return 'Gameover'
        return 'Game'

    def not_game(state):
        game_field.draw(stdscr)
        action = get_user_action(stdscr)
        responses = defaultdict(lambda : state)
        responses['Restart'], responses['Exit'] = 'Init', 'Exit'
        return responses[action]
    state_action = {
        'Init':init,
        'Win':lambda : not_game('Win'),
        'Game':game,
        'Gameover':lambda : not_game('Gameover')
    }

    state = 'Init'
    while state != 'Exit':
        state = state_action[state]()

curses.wrapper(main)


import random
import sys
from PySide6.QtWidgets import QApplication , QMainWindow ,  QMessageBox
from PySide6.QtUiTools import QUiLoader
from ui_guess import Ui_MainWindow
from functools import partial

class Guess_Number(QMainWindow):
    def __init__(self) :
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.guessbtn = self.ui.guessbtn
        self.textbtn = self.ui.textbtn
        self.resultbtn = self.ui.resultbtn
        self.user_num = 0
        self.counter = 0
        self.computer_num = random.randint(1 , 50)
        self.guessbtn.clicked.connect(self.play)

    def play(self):

        self.user_num = int(self.textbtn.toPlainText())
        while self.computer_num != self.user_num  :
            self.counter +=  1
            if self.computer_num > self.user_num :
                self.resultbtn.setText(" GO UPðŸ”ºðŸ”º\n"+"("+str(self.counter)+" guess)")
                break
            elif self.computer_num < self.user_num :
                self.resultbtn.setText(" GO DOWNðŸ”»ðŸ”»\n"+"("+str(self.counter)+" guess)")
                break
        else :
            if self.computer_num == self.user_num :
                message = QMessageBox()
                message.setText("ðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹\nðŸŸ£ðŸ”¹ðŸŸ¢You Won After " + str(self.counter) + " TimesðŸ”¹ðŸŸ£ðŸ”¹\nðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹ðŸŸ£ðŸ”¹ðŸŸ¢ðŸ”¹")
                message.exec()

app = QApplication(sys.argv)
main_window = Guess_Number()
main_window.show()

app.exec()


# -*- coding: utf-8 -*-
"""
Created on Thu Jan 30 13:40:54 2020

@author: engelen

Process MODFLOW files to .nc and initital .IDFs for iMOD-SEAWAT
Only useful for a private branch of IMOD-WQ where .IDFs are read and MODFLOW
files are written, to make up for the >100k of .IDFs written as modeloutput.

"""

import flopy
from itertools import islice
import re
import pandas as pd
import io
from glob import glob
import numpy as np
import xarray as xr
import cftime
from collections import defaultdict
from imod import idf
import os, sys

def idrange2coord(start, stop, dcell):
    """Convert range of ids to coordinates
    """
    return(np.arange(start,stop)*dcell+dcell/2.)

def find_line(pattern, header):
    return([line for line in header if re.match(pattern, line) is not None])

def pattern2df(pattern, path_listfile):
    with open(path_listfile) as listfile:
        header = islice(listfile, 0, 30000)
        lines = find_line(pattern, header)
    return(fwf2df(lines))

def fwf2df(str_list):
    fwf = io.StringIO("".join(str_list))
    return(pd.read_fwf(fwf, header = None))

def combine_all(ds_list, flip_y=True):
    ds_tot = ds_list[0]
    for ds in ds_list[1:]:
        ds_tot = ds_tot.combine_first(ds)

    if flip_y:
        ds_tot.assign_coords(
                y= ds_tot.y.sel(y=slice(None, None, -1))
                )

    return(ds_tot)

def parse_species(paths):
    return([int(re.search("MT3D00(\d).UCN", path).group(1)) for path in paths])

def add_metadata(da, z, dz, dcell):
    """Assign some extra metadata for iMOD-python
    """
    da = da.assign_coords(z=xr.ones_like(da.layer) * z)
    da = da.assign_coords(dz=xr.ones_like(da.layer) * dz)
    da["dx"] = dcell
    da["dy"] = -dcell
    return(da)

def calc_fresh_water_head(head, conc, dense_ref=1000., denselp=0.7143):
    """See Post, Kooi & Simmons (2007)
    Using hydraulic head measurements in variable-density ground water flow analyses
    """
    rho_i = dense_ref+conc*denselp
    return(rho_i/dense_ref * head - (rho_i - dense_ref)/dense_ref * head.z)

#%%Path management
if len(sys.argv) > 1:
    modelfol  = sys.argv[1]
    mod_nr = int(sys.argv[2])
else:
    ##For Testing
    modelfol = r"g:\test_UCN\AD_i153_nr00"
    mod_nr = 0

run_path=os.path.join(modelfol, "*{}.run".format(mod_nr))

r = re.compile("([a-zA-Z]+_i[0-9]+_nr)([0-9]+)")
mname =  r.match(os.path.basename(modelfol)).group(1)

globpath_ucn = os.path.join(modelfol, r"MT3D00[0-9].UCN.p*")
globpath_hds = os.path.join(modelfol, r"*.hds.p*")
path_l = os.path.join(modelfol, "seawat_tmp", r"{}{:02d}.list.p000".format(mname.lower(), mod_nr))
path_dummy_idf = os.path.join(modelfol, "bas", "ibound_l1.idf")

paths_hds = glob(globpath_hds)
paths_hds.sort()

paths_ucn = glob(globpath_ucn)
paths_ucn.sort()

species = parse_species(paths_ucn)

d1 = defaultdict(list)
for k, v in zip(species, paths_ucn):
    d1[k].append(v)

#%%Cellsize
dummy = idf.open(path_dummy_idf)
dcell = dummy["dx"].values

#%%Read UCN data
ucn_data = {}
for s, paths in d1.items():
    ucn_data[s] = dict([(
            i, flopy.utils.binaryfile.UcnFile(path).get_alldata()
            ) for i, path in enumerate(paths)])

#%%Read hds data
hds_data = [flopy.utils.binaryfile.HeadFile(path).get_alldata() for path in paths_hds]

#%%Parse listfile to extract subdomains
df_pks = pattern2df(r" p\d{03} :", path_l)
df_mod = df_pks.loc[:, :8] #Overlap 1 cell for implicit solvers
df_mt  = df_pks.loc[:, 10:] #Overlap 2 cells for advection package

#Convert FORTRAN 1-based indexing to Python 0-based indexing
df_mod.loc[:, [5,6]] -= 1
df_mod.loc[:, [7,8]] -= 1

df_mt.loc[:, [13, 14]] -= 1
df_mt.loc[:, [15, 16]] -= 1

#%%Parse listfile for other coordinates...
df_top = pattern2df(r" TOP ELEVATION OF LAYER 1 =", path_l)
df_btm = pattern2df(r"   MODEL LAYER BOTTOM EL. =", path_l)

topbots = pd.concat([df_top[6], df_btm[5]]).reset_index(drop=True)

z = (topbots[1:].values + topbots[:-1].values)/2
dz=topbots[:-1].values - topbots[1:].values
layer = np.arange(len(z))+1
#%%Calculate times
start_year = 1999
times = flopy.utils.binaryfile.UcnFile(paths_ucn[0]).times
years = start_year + np.round(np.array(times)/365.25, 0)
years = [cftime.DatetimeProlepticGregorian(year, 1, 1) for year in years]

#%%Calculate proxy coordinates
#These are only used with concatenation
x = [idrange2coord(i, j+1, dcell) for i, j in df_mt[[13, 14]].values]
y = [idrange2coord(i, j+1, dcell) for i, j in df_mt[[15, 16]].values]

#%%Create DataArrays and combine these into a big one.
dims = ("time", "layer", "y", "x")

das_s = [(s, [xr.DataArray(ucn, dims = dims,
                         coords={"x" : x[i], "y": y[i], "layer": layer, "time": years}
                         ) for i, ucn in ucn_data[s].items()]) for s in ucn_data.keys()]

das_h = [xr.DataArray(hds, dims = dims,
                      coords={"x" : x[i], "y": y[i], "layer": layer, "time": years}
                      ) for i, hds in enumerate(hds_data)]

da_all = [combine_all(das) for s, das in das_s]
da_all = xr.concat(da_all, dim="species").assign_coords(species = list(ucn_data.keys()))

da_hds = combine_all(das_h)
da_hds = da_hds.where(da_hds<1e20)

#Override coordinates with coordinates dummy idf, because we cannot infer from the UCN and .hds file
#which rows/columns had been removed during dropna when creating the model
da_all = da_all.assign_coords(x=dummy.x, y=dummy.y)
da_hds = da_hds.assign_coords(x=dummy.x, y=dummy.y)

#%%Assign some extra metadata for iMOD-python
da_all = add_metadata(da_all, z, dz, dcell)
da_hds = add_metadata(da_hds, z, dz, dcell)
da_all = da_all.rename("conc")

da_all = xr.where(np.isfinite(da_all), da_all, -9999)
da_hds = da_hds.where(da_hds<1e20, 1e30)

#%%Save inits
shead = da_hds.isel(time=-1)
sconc = da_all.isel(time=-1)

idf.save(os.path.join(modelfol, "..", mname+"{:02d}".format(mod_nr+1), "bas", "head"), shead)
idf.save(os.path.join(modelfol, "..", mname+"{:02d}".format(mod_nr+1), "btn", "conc"), sconc,
         pattern = r"{name}_{time:%Y%m%d%H%M%S}_c{species}_l{layer}{extension}")

#%%Split species into seperate variables
#Can only provide name
ds = da_all.to_dataset(dim="species")
ds = ds.rename(name_dict = dict([(s, "conc{}".format(s)) for s in np.unique(species)]))

#%%Deal with NaNs and assign head to dataset
ds["head"] = da_hds.where(da_hds<1e20)
ds["fhead"] = calc_fresh_water_head(ds["head"], ds["conc1"])

ds["head"]  = xr.where(np.isfinite(ds["head"]),  ds["head"],  1e30)
ds["fhead"] = xr.where(np.isfinite(ds["fhead"]), ds["fhead"], 1e30)

#%%Save to netcdf
ds.time.encoding["units"]    = r"days since 1999-01-01"
ds.time.encoding["calendar"] = r"proleptic_gregorian"

ds = ds.swap_dims({"layer" : "z"})

os.makedirs(os.path.join(modelfol, "results"), exist_ok=True)

ds.to_netcdf(os.path.join(modelfol, "results", "results_{:03d}.nc".format(mod_nr)),
                 unlimited_dims=["time"])


"""Set up of resources for app."""
import sqlite3
from flask_restful import Resource, reqparse
from flask_jwt import jwt_required

class Item(Resource):
    """Resource for item endpoints."""

    parser = reqparse.RequestParser()
    parser.add_argument(
        'price',
        type=float,
        required=True,
        help="This field cannot be left blank!")

    @classmethod
    def find_by_name(cls, name):
        """Find User by name in database and return, if not return none."""
        connection = sqlite3.connect('data.db')
        cursor = connection.cursor()

        query = "SELECT * FROM items WHERE name=?"
        result = cursor.execute(query, (name,))
        row = result.fetchone()
        connection.close()

        if row:
            return {'item': {'name': row[0], 'price': row[1]}}

    @classmethod
    def insert(cls, item):
        """Insert an item into the database."""
        connection = sqlite3.connect('data.db')
        cursor = connection.cursor()

        query = "INSERT INTO items VALUES (?, ?)"
        cursor.execute(query, (item['name'], item['price']))

        connection.commit()
        connection.close()

    @classmethod
    def update(cls, item):
        """Update an item in the database."""
        connection = sqlite3.connect('data.db')
        cursor = connection.cursor()

        query = "UPDATE items SET price=? WHERE name=?"
        cursor.execute(query, (item['price'], item['name']))

        connection.commit()
        connection.close()

    @jwt_required()
    def get(self, name):
        """GET /item route."""
        item = Item.find_by_name(name)
        if item:
            return item
        return {'message': 'Item not found'}, 404

    def post(self, name):
        """POST /item route."""
        if Item.find_by_name(name):
            return {
                'message': f"An item with name '{name}' already exists."
            }, 400

        data = Item.parser.parse_args()

        item = {'name': name, 'price': data['price']}

        try:
            Item.insert(item)
        except:
            return {'message': 'An error occured inserting the item.'}, 500

        return item, 201

    def delete(self, name):
        """DELETE /item route."""
        connection = sqlite3.connect('data.db')
        cursor = connection.cursor()

        query = "DELETE FROM items WHERE name=?"
        cursor.execute(query, (name,))

        connection.commit()
        connection.close()

        return {'message': 'Item deleted'}

    def put(self, name):
        """PUT /item route."""
        data = Item.parser.parse_args()

        item = Item.find_by_name(name)
        updated_item = {'name': name, 'price': data['price']}

        if item is None:
            Item.insert(updated_item)
        else:
            Item.update(updated_item)
        return updated_item

class ItemList(Resource):
    """Resource for items endpoint."""

    def get(self):
        """GET /items route."""
        connection = sqlite3.connect('data.db')
        cursor = connection.cursor()

        query = "SELECT * FROM items"
        result = cursor.execute(query)
        items = []

        for row in result:
            items.append({
                'name': row[0],
                'price': row[1]
            })

        connection.close()

        return {'items': items}


import numpy as np
try:
    from tqdm import tqdm
except:
    tqdm = None

class walker:
    def __init__(self, scale=1, tune_interval=100):
        self.scale = scale
        self.tune_interval = tune_interval
        self.steps_until_tune = tune_interval
        self.accepted = 0
    def step(self, q0):
        if self.steps_until_tune == 0:
            acc_rate = self.accepted / self.tune_interval
            if acc_rate < 0.001:
                self.scale *= 0.1
            elif acc_rate < 0.05:
                self.scale *= 0.5
            elif acc_rate < 0.2:
                self.scale *= 0.9
            elif acc_rate > 0.95:
                self.scale *= 10.0
            elif acc_rate > 0.75:
                self.scale *= 2.0
            elif acc_rate > 0.5:
                self.scale *= 1.1
            self.steps_until_tune = self.tune_interval
            self.accepted = 0
        self.steps_until_tune -= 1
        q = q0 + np.random.normal() * self.scale
        return q
    def accept(self, lnp0, lnp):
        if lnp == -np.inf:
            return False
        if -np.inf < np.log(np.random.uniform()) < lnp - lnp0:
            self.accepted += 1
            return True
        return False

class Sampler:
    def __init__(self, q0, lnprob, args=(), scale=None, tune_interval=None, progress_bar=True):
        self.progress_bar = progress_bar
        self.lnprob = lnprob
        self.args = args
        self.walkers = []
        self.q = np.array(q0, dtype=np.float64)
        self.has_blobs = False
        self.lnp0 = self.lnprob(self.q, *self.args)
        try:
            self.lnp0, blob = self.lnp0
            self.has_blobs = True
        except TypeError:
            pass
        if scale is None:
            scale = [1 for i in self.q]
        if tune_interval is None:
            tune_interval = [100 for i in self.q]
        for i in range(len(self.q)):
            self.walkers.append(walker(scale[i], tune_interval[i]))
    def run(self, nsamples, burnin=0):
        samples = np.empty((nsamples, self.q.size), dtype=np.float64)
        blobs = []
        if tqdm is None or not self.progress_bar:
            iter_samples = range(nsamples + burnin)
        else:
            iter_samples = tqdm(range(nsamples + burnin))
        for i in iter_samples:
            if i < burnin:
                self.sample()
            else:
                if self.has_blobs:
                    samples[i-burnin, :], blob = self.sample()
                    blobs.append(blob)
                else:
                    samples[i-burnin, :] = self.sample()
        if self.has_blobs:
            return samples, blobs
        else:
            return samples
    def sample(self):
        for i, walker in enumerate(self.walkers):
            q0 = self.q[i]
            self.q[i] = walker.step(q0)
            lnp = self.lnprob(self.q, *self.args)
            if self.has_blobs:
                lnp, blob = lnp
            if walker.accept(self.lnp0, lnp):
                self.lnp0 = lnp
            else:
                self.q[i] = q0
        if self.has_blobs:
            return self.q, blob
        else:
            return self.q


import argparse
import json
import os

import librosa
import numpy as np
import matplotlib.pyplot as plt
from scipy.io.wavfile import write

import osc_gen_modif.dsp as dsp
from osc_gen_modif.dsp import slice_cycles
from osc_gen_modif.sig import SigGen

"""
Long grain introduces click at the change of grain.
- use olap ?

"""

class PreprocessorGrains:
    def __init__(self, sampling_rate, window_size, output_dir):
        self.sampling_rate = sampling_rate

        # FFT parameters
        assert ((window_size & (window_size - 1)) == 0), "window_size is not a power of 2"
        self.window_size = window_size  #  Used in fft when training the model later
        self.hop_length = int(window_size // 4)

        #
        self.output_dir = output_dir

        # #  Allowed f0 for grains
        # #  Determined so that the window_size is a multiple of a complete cycle's length
        # self.allowed_f0s = []
        # allowed_f0 = 1
        # while allowed_f0 <= self.window_size:
        #     self.allowed_f0s.append(allowed_f0)
        #     allowed_f0 *= 2

        return

    def preprocess_file(self, filepath):
        print(f'# {filepath}')
        filename = os.path.splitext(os.path.basename(filepath))[0]
        ############################################################
        #  Read file
        y, _ = librosa.core.load(filepath, sr=self.sampling_rate, mono=True)
        #  Remove leading and trailing silences
        y, _ = librosa.effects.trim(y)

        #  Just a test
        write(f'dump/{filename}.wav', self.sampling_rate, y)

        ###########################################################
        # Note: tried osc_gen
        #  but it's a bit too coarse for using.
        #  for instance a unique frequency is used
        # sig_gen = SigGen(num_points=self.window_size)
        # cycles = slice_cycles(y, 50, self.sampling_rate)
        # waves = [sig_gen.arb(c) for c in cycles]

        ###########################################################
        # Detect successive zeros corresponding to cycling points
        #  Instantaneous frequencies
        inst_freq, _ = librosa.core.piptrack(y, sr=self.sampling_rate,
                                             n_fft=self.window_size,
                                             hop_length=self.hop_length)

        #  start_frame is the next zero
        zero_crossings_left = np.where(np.abs(np.diff(np.sign(y))) > 0)[0]

        # Browse waveform, extracting segment of size window_length
        start_frame = 0
        counter = 0
        while start_frame < len(y):
            #  We want after the zero crossing
            index_zero_crossing = np.min(np.where(zero_crossings_left >= start_frame))
            start_frame = zero_crossings_left[index_zero_crossing] + 1

            # Get instantaneous f0
            segment_index = int(start_frame / self.hop_length)
            detected_inst_freq = np.where(inst_freq[:, segment_index] > 0)
            non_zeros_freq = inst_freq[detected_inst_freq, segment_index]
            f0 = np.min(non_zeros_freq)

            ################################################
            # How to compute end_frame ?
            #  Long grains, finding zeros
            # end_frame = self.find_next_stiching_frame(start_frame, self.window_size, zero_crossings_left, f0)

            #  Short grains (one cycle), on zeros ?
            frames_per_cycle = int(self.sampling_rate // f0)
            end_frame = self.find_next_stiching_frame(start_frame, frames_per_cycle, zero_crossings_left, f0)

            #  Short grains, bourrin
            #  In fact, still need to use frames_per_cycles.
            #  If we use an arbitrary value (like 441 ~ 100Hz @ 44100), artifacts are introduced
            # end_frame = start_frame + frames_per_cycle + 1
            ################################################

            # Extract segment
            segment = y[start_frame:end_frame]

            # Interpolate segment
            def arb(data, num_points):
                """ Generate an arbitrary wave cycle. The provided data will be
                interpolated, if possible, to occupy the correct number of samples for
                a single cycle at our reference frequency and then normalized and
                scaled as appropriate.

                @param data seq : A sequence of samples representing a single cycle
                    of a wave
                """

                interp_y = data
                num = interp_y.size
                interp_x = np.linspace(0, num, num=num)
                interp_xx = np.linspace(0, num, num=num_points)
                interp_yy = np.interp(interp_xx, interp_x, interp_y)
                return interp_yy

            interpolated_segment = arb(segment, self.window_size)

            # Normalize waveform
            norm_interp_seg = 0.999969 * librosa.util.normalize(interpolated_segment, axis=0)
            norm_segment = 0.999969 * librosa.util.normalize(segment, axis=0)

            #  Generate a waveform from the grain at different pitches
            #  Here test if it works to use non perfectly cut waveforms
            wave_to_write = np.concatenate([norm_interp_seg for _ in range(1000)])
            write(f'dump/{filename}_{counter}_interp.wav', self.sampling_rate, wave_to_write)
            #  If we repitch shift the low frequency grain, what happens ?
            shifted_norm_interp_seg = librosa.effects.pitch_shift(wave_to_write, sr=self.sampling_rate, n_steps=42)
            write(f'dump/{filename}_{counter}_pitchShifted.wav', self.sampling_rate, shifted_norm_interp_seg)
            # short grains
            wave_to_write = np.concatenate([norm_segment for _ in range(1000)])
            write(f'dump/{filename}_{counter}.wav', self.sampling_rate, wave_to_write)

            #  Compute next starting frame
            start_frame = self.find_next_stiching_frame(start_frame, self.hop_length, zero_crossings_left, f0)
            counter += 1

        return

    def find_next_stiching_frame(self, start_frame, target_length, zero_crossings_left, f0):
        # Find next stiching point (possibly not at self.window_size frames from start point)
        # which is close to a multiple of Tn0 = fs / f0
        frames_per_cycle = self.sampling_rate // f0
        num_cycles = target_length // frames_per_cycle
        end_frame_approx = start_frame + num_cycles * frames_per_cycle
        #  Take next zero_crossing frame
        end_frame = zero_crossings_left[np.where(end_frame_approx <= zero_crossings_left)[0][0]] + 1
        return end_frame

    def process_database(self, filename):
        with open(filename, encoding='utf-8') as f:
            files = f.readlines()

        for file in files:
            file = file.rstrip()
            self.preprocess_file(file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--filename', type=str,
                        help='txt list containing the files to process')
    parser.add_argument('-c', '--config', type=str,
                        help='JSON config')
    parser.add_argument('-o', '--output_dir', type=str,
                        help='path to output directory')
    args = parser.parse_args()

    # FFT parameters
    with open(args.config) as f:
        data = f.read()
    config = json.loads(data)
    data_config = config["data_config"]
    window_size = data_config["segment_length"]
    sampling_rate = data_config["sampling_rate"]

    preprocessor = PreprocessorGrains(sampling_rate=sampling_rate, window_size=window_size, output_dir=args.output_dir)
    preprocessor.process_database(args.filename)


import datetime
import glob
import time
import h5py
import numpy as np
import matplotlib.pyplot as plt

import PlotFunctions as Plt
import Waveform as Wvf

class Dataset:
    def __init__(self,  Path, ShowPlots=True, Selection='*', Pol=1, NumChannels=2):
        self.Path = Path
        self.NumChannels = NumChannels
        self.ShowPlots = ShowPlots
        self.Selection = Selection
        self.Pol = Pol
        self.Ch = self.InitializeChannels(self.NumChannels, self.Pol)
        self.Files = glob.glob(self.Path+self.Selection)

    def RunStandardAnalysis(self, NoiseDataset=None):
        # self.Files = self.ChooseFilesToAnalyze(self.Path)
        for File in self.Files:
            self.ImportDataFromHDF5(File, self.Ch)
        self.DoAnalysis(self.Ch, NoiseDataset=NoiseDataset)
        self.ChargeCollection = self.Ch[0].Max / self.Ch[1].Max
        self.DiffMinute = int((np.max(self.Ch[0].TimeStamp) - np.min(self.Ch[0].TimeStamp)).seconds/60.0 + 0.5)
        self.XTicks = int((self.DiffMinute/12.0 + 0.5))+1
        self.NoiseCut = 1000
        self.Cut = np.where(self.Ch[0].BaseStd < self.NoiseCut)
        self.InverseCut = np.where(self.Ch[0].BaseStd > self.NoiseCut)

    def InitializeChannels(self, NumChannels=2, Pol=1):
        return [Wvf.Waveform(ID=ii, Pol=(-1)**ii*-1*Pol) for ii in range(1,NumChannels+1)]

    # def ChooseFilesToAnalyze(self, Path):
    #     return

    def ImportDataFromHDF5(self, File, channels):
        f = h5py.File(File, 'r')
        # print(" | Filename...", File)
        Keys = list(f.keys())
        for ch in channels:
            ch.Time = np.array(f.get('Time')).flatten() * ch.TScale
            ch.Trigger = np.array(f.get('Trigger')).flatten() * ch.VScale
            Group = f.get(ch.ChName)
            GroupKeys = Group.keys()
            ch.Files.append(len(GroupKeys))
            # print(" | Number of files in ch%d...\t %d/%d" % (ch.ID, ch.Files[-1], np.sum(ch.Files)))
            for key in GroupKeys:
                ch.Amp.append(np.array(Group.get(key)).flatten() * ch.VScale * ch.Pol)
                ch.TimeStamp.append(datetime.datetime.strptime((f.attrs['Date']+Group.get(key).attrs["TimeStamp"]), '%Y%m%d%H%M%S'))
        f.close()

    def DoAnalysis(self, channels, NoiseDataset=None):
    ###### Basic analysis: baseline subtraction, waveform averaging, obtaining fourier spectra, frequency bandpass filter and finding extrema.
        Print = False
        for ii, ch in enumerate(channels):
            # print(" | Processing data in channel %d..." % (ch.ID))
            ch.GetSampling()
            ch.Amp = [x for _, x in sorted(zip(ch.TimeStamp, ch.Amp))]
            ch.Amp = np.array(ch.Amp)
            ch.TimeStamp = np.array(sorted(ch.TimeStamp))

            # ch.TimeStamp = np.array(ch.TimeStamp)
            ch.Amp = ch.SubtractBaseline(Data=ch.Amp, state=Print)
            ch.Amp = ch.RemoveNoise(Data=ch.Amp, HighPass=80000, state=Print)
            if NoiseDataset is not None:
                for jj,amp in enumerate(ch.Amp):
                    ch.Amp[jj] =  ch.Amp[jj]-np.mean(NoiseDataset.Ch[ii].Amp,axis=0)

            # ch.RunFit(Data=ch.Amp)
            ch.GetAllMaxima(Data=ch.Amp, state=Print)
            # ch.FindMaxGradient(Data=ch.Amp ,state=Print)
            ch.GetDriftTime(Data=ch.Amp, Threshold=0.1)
            ch.GetIntegral(Data=ch.Amp, state=Print)
            # ch.GetBaselineNoise(Data=ch.Amp)

        # print(" | Time elapsed: ", time.process_time() , "sec")

    def ShowStandardPlots(self):
        self.ShowBaselineNoise()
        self.ShowAmplitudeVsTime()
        self.ShowChargeCollection()
        self.ShowDrifttimeVsTime()

    def ShowBaselineNoise(self, Channel=-1, BinMax=20):
        fig = plt.figure(figsize=(8,6))
        plt.xlim(0,BinMax)
        plt.xlabel('Baseline RMS Noise [mV]')
        plt.ylabel('Counts/bin')
        histmax = 0
        for ii in range(self.NumChannels):
            if ii+1 is not Channel and Channel != -1:
                continue

            h,hx,hp = plt.hist(self.Ch[ii].BaseStd[self.Cut], bins=np.arange(0.0,BinMax,0.2), histtype='step', align='mid', lw=2, color=Plt.colors[ii], label=self.Ch[ii].Name)
            plt.axvline(np.median(self.Ch[ii].BaseStd[self.Cut]), color=Plt.colors[ii])
            rectangle = plt.Rectangle(xy=(np.median(self.Ch[ii].BaseStd[self.Cut])-np.std(self.Ch[ii].BaseStd[self.Cut])/np.sqrt(len(self.Ch[ii].BaseStd[self.Cut])),0),
                                    width=2*np.std(self.Ch[ii].BaseStd[self.Cut])/np.sqrt(len(self.Ch[ii].BaseStd[self.Cut])),
                                    height=10000,
                                    fc=Plt.colors[ii],
                                    ec=Plt.colors[ii],
                                    alpha=0.05,
                                    fill=True)
            plt.gca().add_patch(rectangle)
            if np.max(h) > histmax:
                histmax = np.max(h)
        plt.ylim(0, self.RoundUpToNext(histmax, 10))
        plt.legend(loc='upper right')
        fig.tight_layout()

    def ShowAmplitudeVsTime(self, Channel=-1, YTicks=100, YMax=None):
        if YMax is None:
            YMax = np.max([np.max(self.Ch[ii].Max) for ii in range(self.NumChannels)])
        Plt.PltChargeVsTime(Time=self.Ch[0].TimeStamp[self.Cut],
                    Data=[self.Ch[0].Max[self.Cut], self.Ch[1].Max[self.Cut]],
                    CC=self.ChargeCollection[self.Cut]*100,
                    Legend=['Anode','Cathode','Charge Collection'],
                    Label='Amplitude [mV]',
                    XTicks=self.XTicks,
                    YTicks=YTicks,
                    YRange=[0,YMax],
                    Title='')

    def ShowIntegralVsTime(self, Channel=-1, YTicks=100, YMax=None):

        self.Ch[0].Integral /= np.max(self.Ch[0].Integral)
        self.Ch[1].Integral /= np.max(self.Ch[1].Integral)
        if YMax is None:
            YMax = np.max([np.max(self.Ch[ii].Integral) for ii in range(self.NumChannels)])
        Plt.PltChargeVsTime(Time=self.Ch[0].TimeStamp[self.Cut],
                    Data=[self.Ch[0].Integral[self.Cut], self.Ch[1].Integral[self.Cut]],
                    CC=self.Ch[0].Integral[self.Cut]/self.Ch[1].Integral[self.Cut]*100,
                    Legend=['Anode','Cathode','Charge Collection'],
                    Label='Integral [a.u.]',
                    XTicks=self.XTicks,
                    YTicks=YTicks,
                    YRange=[0,YMax],
                    Title='')

    def ShowDrifttimeVsTime(self, YMax=0, Channel=-1):
        # YMax = np.max([np.max(self.Ch[ii].DriftTime) for ii in range(self.NumChannels)])
        YMax = self.RoundUpToNext(YMax, 10)
        YTicks = self.RoundDownToNext(YMax/5, 1)
        self.DriftTime = self.Ch[0].DriftTime - self.Ch[1].DriftTime
        DriftCut = np.where(self.DriftTime > 0)
        Plt.PltTime(Time=self.Ch[0].TimeStamp[DriftCut],
                    Data=[self.Ch[0].DriftTime[DriftCut], self.Ch[1].DriftTime[DriftCut], self.Ch[0].DriftTime[DriftCut] - self.Ch[1].DriftTime[DriftCut]],
                    Legend=['Anode','Cathode','Drift Time'],
                    Label='Time since Trigger [$\mu$s]',
                    XTicks=self.XTicks,
                    YTicks=YTicks,
                    YRange=[0,YMax],
                    SaveName='drift_time',
                    Save=False)

    def ShowChargeCollection(self, Channel=-1):
        Mean = np.mean(self.ChargeCollection)
        Median = np.median(self.ChargeCollection)
        Err = np.std(self.ChargeCollection)/np.sqrt(len(self.ChargeCollection))

        XMin = self.RoundDownToNext(np.min(self.ChargeCollection), 1)
        XMax = self.RoundUpToNext(np.max(self.ChargeCollection), 1)
        fig = plt.figure(figsize=(8,6))

        h,hx,hp = plt.hist(self.ChargeCollection, bins=np.arange(XMin,XMax,0.01), histtype='step', align='mid', lw=2)

        plt.xlim(XMin, self.RoundUpToNext(np.max(self.ChargeCollection), 1))
        plt.ylim(0, self.RoundUpToNext(np.max(h), 10))

        plt.xlabel('Charge Collection')
        plt.ylabel('Counts/bin')
        plt.axvline(Mean, color='red', )
        plt.title('Charge Collection = %.4f ± %.4f' % (Mean, Err))
        rectangle = plt.Rectangle((Mean-Err,0), 2*Err, 10000, fc='red',ec="red", alpha=0.2, fill=True,)
        plt.gca().add_patch(rectangle)
        fig.tight_layout()

    def RoundUpToNext(self, Num, Ceil):
        return int(np.ceil(Num / float(Ceil))) * float(Ceil)

    def RoundDownToNext(self, Num, Floor):
        return int(np.floor(Num / float(Floor))) * float(Floor)


import os
import numpy as np
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras import models, layers, activations, losses, optimizers, metrics
import callback_helper

imbd_dir = '/home/jp/workspace/aclImdb'
train_dir = os.path.join(imbd_dir, 'train')
test_dir = os.path.join(imbd_dir, 'test')

def load_data(text_dir):
    texts = []
    labels = []

    for label_type in ['neg', 'pos']:
        dir_name = os.path.join(text_dir, label_type)
        for fname in os.listdir(dir_name):
            if fname[-4:] == '.txt':
                with open(os.path.join(dir_name, fname)) as f:
                    texts.append(f.read())
                if label_type == 'neg':
                    labels.append(0)
                else:
                    labels.append(1)

    return np.asarray(texts), np.asarray(labels)

texts, labels = load_data(train_dir)

maxlen = 100
train_samples = 200
validation_samples = 10000
max_words = 10000

tokenizer = Tokenizer(num_words=max_words)
tokenizer.fit_on_texts(texts)
word_index = tokenizer.word_index

print('Found %s unique tokens' % len(word_index))

sequences = tokenizer.texts_to_sequences(texts)

data = pad_sequences(sequences,
                     padding='post',
                     truncating='post',
                     maxlen=maxlen)
labels = np.asarray(labels)

indices = np.arange(data.shape[0])
np.random.shuffle(indices)

# 这种方式的赋值，是numpy功能，普通list无效
data = data[indices]
labels = labels[indices]

# 将数据分成训练和检验两部分
x_train = data[:train_samples]
y_train = labels[:train_samples]
x_validation = data[train_samples:train_samples + validation_samples]
y_validation = labels[train_samples:train_samples + validation_samples]

# # 解析GloVe词嵌入文件
glove_dir = '/home/jp/workspace/glove.6B'
embeddings_index = {}
with open(os.path.join(glove_dir, 'glove.6B.100d.txt')) as f:
    for line in f:
        values = line.split()
        word = values[0]
        coefs = np.asarray(values[1:], dtype='float32')
        embeddings_index[word] = coefs

print('Found %s words in vectors.' % len(embeddings_index))

# # 准备GloVe词嵌入矩阵
embeddings_dim = 100
embeddings_matrix = np.zeros((max_words, embeddings_dim))
for word, index in word_index.items():
    if index < max_words:
        embeddings_vector = embeddings_index.get(word)
        if embeddings_vector is not None:
            embeddings_matrix[index] = embeddings_vector

# 定义模型
model = models.Sequential()
model.add(layers.Embedding(max_words, 100, input_length=maxlen))
model.add(layers.Flatten())
model.add(layers.Dense(21, activation=activations.relu))
model.add(layers.Dense(1, activation=activations.sigmoid))
model.summary()

# 将预训练的词嵌入加载到Embedding层
model.layers[0].set_weights([embeddings_matrix])
model.layers[0].trainable = False

model.compile(optimizer=optimizers.rmsprop(),
              loss=losses.binary_crossentropy,
              metrics=['acc'])

callbacks = callback_helper.generate_callbacks()

history = model.fit(x_train,
                    y_train,
                    epochs=10,
                    batch_size=32,
                    validation_data=(x_validation, y_validation),
                    callbacks=callbacks)

# 准备测试数据
x_test, y_test = load_data(test_dir)
x_test = tokenizer.texts_to_sequences(x_test)
x_test = pad_sequences(x_test,
                       padding='post',
                       truncating='post',
                       maxlen=maxlen)

# 在测试集上测试
result = model.evaluate(x_test, y_test)
print(result)


import numpy as np
import torch

from .base_selector import BaseSelector
from ...config.setting.core_setting import (
    INDICES_TUPLE,
    INDICES_FLAG
)

class DistanceWeightedSelector(BaseSelector):
    """
    Distance weighted sampling method, euclidean distance metric is required.

    paper: `Sampling Matters in Deep Embedding Learning <http://openaccess.thecvf.com/content_iccv_2017/html/Wu_Sampling_Matters_in_ICCV_2017_paper.html>`_
    """
    def __init__(self, lower_cutoff=0.5, upper_cutoff=1.4, embedding_dim=512, **kwargs):
        super(DistanceWeightedSelector, self).__init__(**kwargs)
        self.lower_cutoff = lower_cutoff
        self.upper_cutoff = upper_cutoff
        self.embedding_dim = embedding_dim

    def forward(
        self,
        metric_mat,
        row_labels,
        col_labels,
        is_same_source=False
    ) -> tuple:
        """
        Randomly select a positive sample for anchor sample and select a negative sample for anchor sample according to the distance weighted probability.

        The distribution of pairwise distances follows:

        :math:`q(d) \propto d^{n-2} [1 - \\frac{1}{4} d^2 ]^{\\frac{n-3}{2}}`
        """
        metric_mat_to_weights = metric_mat.clamp(min=self.lower_cutoff).detach()
        bs = metric_mat.size(0)
        device = metric_mat.device
        # pos and neg mask
        matches = (row_labels == col_labels).byte()
        diffs = matches ^ 1

        # compute distance weighted prob
        minus_log_weights = (
            (2.0 - self.embedding_dim) * torch.log(metric_mat_to_weights) +
            ((3.0 - self.embedding_dim) / 2.0) * torch.log(1.0 - 0.25 * metric_mat_to_weights.pow(2))
        ) * diffs
        ## avoid inf or nan
        inf_or_nan = torch.isinf(minus_log_weights) | torch.isnan(minus_log_weights)
        ## compute weight
        inv_weights = torch.exp(
            minus_log_weights - torch.max(minus_log_weights[~inf_or_nan])
        )
        ## NOTE: Cutting of values with high distances made the results slightly worse according to Confusezius / Pytorch-Metric-Learning-Baseline
        # inv_weights = (
        #     inv_weights * diffs * (metric_mat_to_weights < self.upper_cutoff).byte()
        # )
        inv_weights[inf_or_nan] = 0

        # construct the tuple
        if is_same_source:
            matches.fill_diagonal_(0)

        # avoid nan, inf and zero
        non_zero_row = torch.where(
            (torch.sum(matches, dim=-1) > 0) & (torch.sum(inv_weights, dim=-1) > 0)
        )[0]
        inv_weights = inv_weights[non_zero_row, :]
        ## anchor indices
        a_ids = torch.arange(bs)[non_zero_row].to(device)
        ## positive indices
        p_ids = torch.multinomial(
            input=matches[non_zero_row, :].float(),
            num_samples=1,
            replacement=True
        ).flatten()
        ## negative indices
        inv_weights = inv_weights / inv_weights.sum(dim=-1, keepdim=True)
        n_ids = torch.multinomial(
            input=inv_weights,
            num_samples=1,
            replacement=True
        ).flatten()

        indices_tuple = {
            INDICES_TUPLE: torch.stack([a_ids, p_ids, n_ids], dim=1),
            INDICES_FLAG: None
        }
        weight = None
        return (
            metric_mat,
            row_labels,
            col_labels,
            is_same_source,
            indices_tuple,
            weight,
        )