import sys
sys.path.append('../rxnft_vae')
import rdkit
import rdkit.Chem as Chem
from rdkit.Chem import QED, Descriptors, rdmolops

import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader
from torch.autograd import Variable

import math, random, sys
from optparse import OptionParser
from collections import deque

from reaction_utils import get_mol_from_smiles, get_smiles_from_mol,read_multistep_rxns, get_template_order, get_qed_score,get_clogp_score
from reaction import ReactionTree, extract_starting_reactants, StartingReactants, Templates, extract_templates,stats
from fragment import FragmentVocab, FragmentTree, FragmentNode, can_be_decomposed
from vae import FTRXNVAE, set_batch_nodeID
from mpn import MPN,PP,Discriminator
from evaluate import Evaluator
import random
import numpy as np
import networkx as nx

from sparse_gp import SparseGP
import scipy.stats as sps
import sascorer

def decode_many_times(model, latent):
	prob_decode = True
	latent_size = model.latent_size
	ft_mean = latent[:, :latent_size]
	rxn_mean = latent[:, latent_size:]
	product_list=[]
	for i in range(50):
		generated_tree = model.fragment_decoder.decode(ft_mean, prob_decode)
		g_encoder_output, g_root_vec = model.fragment_encoder([generated_tree])
		product, reactions = model.rxn_decoder.decode(rxn_mean, g_encoder_output, prob_decode)
		if product != None:
			product_list.append([product, reactions])
	if len(product_list) == 0:
		return None
	else:
		return product_list

def run_bo(X_train, y_train, X_test, y_test, model, parameters, metric, randseed):
	random_seed = int(randseed)
	np.random.seed(random_seed)
	if metric =="logp":
		logp_m = parameters[0]
		logp_s = parameters[1]
		sascore_m = parameters[2]
		sascore_s = parameters[3]
		cycle_m = parameters[4]
		cycle_s = parameters[5]

	filename = "../Results/" + metric + str(random_seed) + ".txt"

	#print("maxmimum score :", np.min(y_train), X_train.shape)
	#print(y_train)
	with open(filename, "w") as writer:
		iteration = 0
		latents = []
		min_scores = []
		while iteration < 5:
			# fit the GP
			#print("maxmimum score :", np.min(y_train), X_train.shape)
			print(iteration)
			np.random.seed(iteration * random_seed)
			M = 500
			sgp = SparseGP(X_train, 0 * X_train, y_train, M)
			sgp.train_via_ADAM(X_train, 0 * X_train, y_train, X_test, X_test * 0, y_test, minibatch_size = 10 * M, max_iterations = 100, learning_rate = 0.001)

			pred, uncert = sgp.predict(X_test, 0 * X_test)
			error = np.sqrt(np.mean((pred - y_test)**2))
			testll = np.mean(sps.norm.logpdf(pred - y_test, scale = np.sqrt(uncert)))
			print('Test RMSE: ', error, ' Test ll: ', testll)

			pred, uncert = sgp.predict(X_train, 0 * X_train)
			error = np.sqrt(np.mean((pred - y_train)**2))
			trainll = np.mean(sps.norm.logpdf(pred - y_train, scale = np.sqrt(uncert)))
			print( 'Train RMSE: ', error, 'Train ll: ', trainll)
			#print( 'Train ll: ', trainll)

			next_inputs, values = sgp.batched_greedy_ei(60, np.min(X_train, 0), np.max(X_train, 0))
			valid_smiles =[]
			new_features =[]
			full_rxn_strs=[]
			values = values.flatten()
			for i in range(60):
				#print(i)
				latent = next_inputs[i].reshape((1,-1))
				#res = model.decode_many_times(torch.from_numpy(latent).float(), 50)
				res= decode_many_times(model, torch.from_numpy(latent).float())
				if res is not None:
					smiles_list = [re[0] for re in res]
					n_reactions = [len(re[1].split(" ")) for re in res]
					#print(n_reactions)
					for re in res:
						smiles = re[0]
						if len(re[1].split(" ")) > 0 and smiles not in valid_smiles:
							#print(smiles, re[1].split(" "))
							valid_smiles.append(smiles)
							new_features.append(latent)
							full_rxn_strs.append(re[1])
				#print(i, res)

			#new_features = np.vstack(new_features)
			scores =[]
			b_valid_smiles=[]
			b_full_rxn_strs=[]
			b_scores=[]
			b_new_features=[]
			for i in range(len(valid_smiles)):
				if metric =="logp":
					mol = rdkit.Chem.MolFromSmiles(valid_smiles[i])
					if mol is None:
						continue
					current_log_P_value = Descriptors.MolLogP(mol)
					current_SA_score = -sascorer.calculateScore(mol)
					cycle_list = nx.cycle_basis(nx.Graph(rdmolops.GetAdjacencyMatrix(mol)))
					if len(cycle_list) == 0:
						cycle_length = 0
					else:
						cycle_length = max([ len(j) for j in cycle_list ])
					if cycle_length <= 6:
						cycle_length = 0
					else:
						cycle_length = cycle_length - 6
					current_cycle_score = -cycle_length
					current_SA_score_normalized = (current_SA_score - sascore_m) / sascore_s
					current_log_P_value_normalized = (current_log_P_value - logp_m) / logp_s
					current_cycle_score_normalized = (current_cycle_score - cycle_m) / cycle_s
					score = current_SA_score_normalized + current_log_P_value_normalized + current_cycle_score_normalized
					scores.append(-score)
					b_valid_smiles.append(valid_smiles[i])
					b_full_rxn_strs.append(full_rxn_strs[i])
					b_new_features.append(new_features[i])
				if metric=="qed":
					mol = rdkit.Chem.MolFromSmiles(valid_smiles[i])
					if mol!=None:
						score = QED.qed(mol)
						scores.append(-score)
						b_valid_smiles.append(valid_smiles[i])
						b_full_rxn_strs.append(full_rxn_strs[i])
						b_new_features.append(new_features[i])
			new_features = np.vstack(b_new_features)
			if len(new_features) > 0:
				X_train = np.concatenate([ X_train, new_features ], 0)
				y_train = np.concatenate([ y_train, np.array(scores)[ :, None ] ], 0)
			iteration+=1

			for i in range(len(b_valid_smiles)):
				line = " ".join([b_valid_smiles[i], b_full_rxn_strs[i], str(scores[i])])
				writer.write(line + "\n")
			#print(iteration, min(scores))

parser = OptionParser()
parser.add_option("-w", "--hidden", dest="hidden_size", default=200)
parser.add_option("-l", "--latent", dest="latent_size", default=50)
parser.add_option("-d", "--depth", dest="depth", default=2)
parser.add_option("-s", "--save_dir", dest="save_path")
parser.add_option("-t", "--data_path", dest="data_path")
parser.add_option("-v", "--vocab_path", dest="vocab_path")
parser.add_option("-m", "--metric", dest="metric")
parser.add_option("-r", "--seed", dest="seed", default=1)
opts, _ = parser.parse_args()

# get parameters
hidden_size = int(opts.hidden_size)
latent_size = int(opts.latent_size)
depth = int(opts.depth)
vocab_path = opts.vocab_path
data_filename = opts.data_path
w_save_path = opts.save_path
metric = opts.metric
seed = int(opts.seed)

# load model
if torch.cuda.is_available():
	#device = torch.device("cuda:1")
	device = torch.device("cuda")
	torch.cuda.set_device(1)
else:
	device = torch.device("cpu")

print("hidden size:", hidden_size, "latent_size:", latent_size, "depth:", depth)
print("loading data.....")
data_filename = opts.data_path
routes, scores = read_multistep_rxns(data_filename)
rxn_trees = [ReactionTree(route) for route in routes]
molecules = [rxn_tree.molecule_nodes[0].smiles for rxn_tree in rxn_trees]
reactants = extract_starting_reactants(rxn_trees)
templates, n_reacts = extract_templates(rxn_trees)
reactantDic = StartingReactants(reactants)
templateDic = Templates(templates, n_reacts)

print("size of reactant dic:", reactantDic.size())
print("size of template dic:", templateDic.size())

n_pairs = len(routes)
ind_list = [i for i in range(n_pairs)]
fgm_trees = [FragmentTree(rxn_trees[i].molecule_nodes[0].smiles) for i in ind_list]
rxn_trees = [rxn_trees[i] for i in ind_list]
data_pairs=[]
for fgm_tree, rxn_tree in zip(fgm_trees, rxn_trees):
	data_pairs.append((fgm_tree, rxn_tree))
cset=set()
for fgm_tree in fgm_trees:
	for node in fgm_tree.nodes:
		cset.add(node.smiles)
cset = list(cset)
if vocab_path is None:
	fragmentDic = FragmentVocab(cset)
else:
	fragmentDic = FragmentVocab(cset, filename =vocab_path)

print("size of fragment dic:", fragmentDic.size())

# loading model

mpn = MPN(hidden_size, depth)
model = FTRXNVAE(fragmentDic, reactantDic, templateDic, hidden_size, latent_size, depth, fragment_embedding=None, reactant_embedding=None, template_embedding=None)
checkpoint = torch.load(w_save_path, map_location=device)
model.load_state_dict(checkpoint)
print("finished loading model...")

print("number of samples:", len(data_pairs))
latent_list=[]
score_list=[]
print("num of samples:", len(rxn_trees))
latent_list =[]
score_list=[]
if metric =="qed":
	for i, data_pair in enumerate(data_pairs):
		latent = model.encode([data_pair])
		#print(i, latent.size(), latent)
		latent_list.append(latent[0])
		rxn_tree = data_pair[1]
		smiles = rxn_tree.molecule_nodes[0].smiles
		score_list.append(get_qed_score(smiles))
if metric =="logp":
	logP_values = np.loadtxt('logP_values.txt')
	SA_scores = np.loadtxt('SA_scores.txt')
	cycle_scores = np.loadtxt('cycle_scores.txt')

	logp_m = np.mean(logP_values)
	logp_s = np.std(logP_values)

	sascore_m = np.mean(SA_scores)
	sascore_s = np.std(SA_scores)

	cycle_m = np.mean(cycle_scores)
	cycle_s = np.std(cycle_scores)
	for i, data_pair in enumerate(data_pairs):
		latent = model.encode([data_pair])
		latent_list.append(latent[0])
		rxn_tree = data_pair[1]
		smiles = rxn_tree.molecule_nodes[0].smiles
		score_list.append(get_clogp_score(smiles, logp_m, logp_s, sascore_m, sascore_s, cycle_m, cycle_s))
latents = torch.stack(latent_list, dim=0)
scores = np.array(score_list)
scores = scores.reshape((-1,1))
latents = latents.detach().numpy()
n = latents.shape[0]
permutation = np.random.choice(n, n, replace = False)
X_train = latents[ permutation, : ][ 0 : np.int(np.round(0.9 * n)), : ]
X_test = latents[ permutation, : ][ np.int(np.round(0.9 * n)) :, : ]
y_train = -scores[ permutation ][ 0 : np.int(np.round(0.9 * n)) ]
y_test = -scores[ permutation ][ np.int(np.round(0.9 * n)) : ]
print(X_train.shape, X_test.shape)
if metric == "logp":
	parameters = [logp_m, logp_s, sascore_m, sascore_s, cycle_m, cycle_s]
else:
	parameters =[]

run_bo(X_train, y_train, X_test, y_test, model, parameters, metric, seed)

import numpy,import scipy,import networkx


import pickle
import os
import time
import numpy as np
import sys
import shutil
from openmmlib import openmmlib
from openmmlib import polymerutils
from openmmlib.polymerutils import scanBlocks
from openmmlib.openmmlib import Simulation
from openmmlib.polymerutils import grow_rw
import pyximport; pyximport.install()
from smcTranslocator_MovingBarrier import smcTranslocatorDirectional
import tools
import random

# -------defining parameters----------
# -- basic loop extrusion parameters--

logname="log.txt"
GPU = 0
LIFETIME = 100   # Processivity of cohesin, Default: 100 for WT, 1000 for Wapl KO
SEPARATION = 200 # Separation LEFs in number of monomers, Default: 200 for WT, 100 for Wapl KO, assuming a  monomer size of 1kb
N = 10000                     # System size in number of monomers
smcStepsPerBlock = 1          # Number of LEF steps between blocks of polymer simulations, Default: 1
stiff = 1                     # Polymer siffness in unit of bead size, default: 1
dens = 0.2                    # density in beads / volume. The density can roughly be estimated by looking at the amount of DNA in a nuclear volume, Default: 0.2
box = (N / dens) ** 0.33      # Define size of the bounding box for Periodic Boundary Conditions, Default: 0.33
data = polymerutils.grow_rw(N, int(box) - 2)  # creates a compact conformation
block = 0  # starting block
stg = 0.8# same as Banigan, van den Berg, Brandao eLife 2020 #0.1  #stall probability at ctcf sites
unstallLEFRate = 0.005 # about once per two typical LEF lifetimes.
ctcf_interval = 300 # 300 kb, e.g., see Busslinger et al. Nature 2017

# -- LEF and transcription dynamics --

# Note that the each rate can have a maximum value of 1.
lef_speed=1.0
lefperm = 0. # controls amount of LEF-LEF bypassing we have
pauseArray = np.ones(N,dtype=np.double) # The speed of LEFS at each position. Default is an array of ones, which means that LEFS go on maximum speed everywhere. To simulate the presence of CTCF Aafke reduced the speed at CTCF sites to 0.005. However, this doesn't account for directionality of CTCF!
shrinkPauseArray=np.zeros(N,dtype=np.double)# speed at which loops shrink
shrink_speed=0.0
kinPol= 0.001     # PolII initiation rate, Default range: 0.00025-0.002
kterPol=0.002                 # Termination rate of PolII, Default range: 0.002-1.0. Normally transcription is initiation limited, so make sure that kinPol < kterPol
kterPolArray=np.zeros(N,dtype=np.double) # fix this later based on gene stucture
PolSpeed=0.1                  # The speed of PolII as a fraction of the speed of cohesin. Default value: 0.1
poldissoc=0.
#pauseArrayPol=np.zeros(N,dtype=np.double) + PolSpeed # The speed of PolII at each lattice site #initialize later
polloading=np.array([950,1950,2950,3950,4950,5950,6950,7950,8950]) # PolII initiation sites. The direction of the gene is set by the relative position of the polloading and poltermination sites.
#np.array([700,903,1100,1303,1700,1903,2100,2303,2700,2903,3100,3303,3700,3903,4100,4303,4700,4903,5100,5303,5700,5903,6100,6303])
poltermination=np.array([1850,2850,3850,4850,5850,6850,7850,8850,9850]) #PolII termination sites. At a termination site, PolII stalls with probability 1 and then unloads with rate kterPol. To simulate a broad PolII unloading area, I would set the poltermination site far beyond the gene length and then define wide region where PolII can stall, as shown on the lines below
#np.array([800,803,1200,1203,1800,1803,2200,2203,2800,2803,3200,3203,3800,3803,4200,4203,4800,4803,5200,5203,5800,5803,6200,6203])
stalProbPol=np.zeros(N,dtype=np.double) # The rate of PolII stalling. Once stalled, PolII unloads with a rate 'kterPol'. One can choose a single stall site, or a range of stall sites. If this array is set to zero everywhere, PolII will stall at the defined termination sites.
stalProb=0.001
stall_in_gene=0.
unstall_in_gene=1.
unstallArray=np.zeros(N,dtype=np.double) #array for unstalling in of Pol II in gene

STALL_FROM_FILE = False # option to take stall probabilities from file for spatially varying patterns
stallfile=""
gene_stall=[]
ctcf_left_list=[]
ctcf_right_list=[]
strongCTCFstall = 0

TSSloadbias = 1.0
TSSloadstart=0 #use this and TSSloadend to control offset/width of loading near TSS,  note: positive TSSloadstart denotes num sites before TSS, positive TSSloadend denotes sites after TSS
TSSloadend=0

base_genelen=200
#genelength=[base_genelen]*len(polloading)
variable_genelength = 0
fixed_variable_genelength=0
variable_permeability = 0 #This varaible toggles on/off variable permeability of cohesin through RNAp depending on position in gene, 0 or 1 for constant permeability or variable perm, respectively
variable_type = 0 #type of variable permeability: 0- linearly decreasing, 1- step function
variable_offset = 0 #param for var perm function: for linear this is a constant offset from 0, for step this is offset for the bottom step
variable_pos = base_genelen // 2 #param: irrelevant for linear, position of step for step function.
variable_return = 0 #return to max permeability after passing end of gene but before reaching termination site
variable_TSSfactor = 1.0  # factor by which permeability at TSS is lower (or higher)

permLeftArray=np.zeros(N,dtype=np.double)
permRightArray=np.zeros(N,dtype=np.double)

collisionLifeFactor=1.0 # factor by which head-on collision with RNAP changes lifetime of cohesin
collisionLife=LIFETIME*collisionLifeFactor
TTSunload=1.0 # factor by which cohesin life time is changed near TTS
unloadZone=10 # width of zone in which cohesin life is changed by TTS

PolPause=0.002   # Step rate PolII at TSS, Default: 0.002
#later:
#for i in polloading:
#    pauseArrayPol[i]=PolPause

L=0.          # Set permeability of PolII to cohesin. L=0 means PolII is impermeable.
R=0.          # Set permeability of PolII to cohesin coming from the right

run_id=1

# -- polymer simulation settings --

steps = int(200*(smcStepsPerBlock)) # nr of 3D simulation blocks btw advancing LEFs. For deterministic stepping choose 200-250 steps per block, otherwise, rescale with stepping probability. When genes are sparse, smcStepsPerBLock is approximately the number of smc steps per smc block.

saveRNAP=True # whether or not to print RNAP positions w/ each printed block

saveEveryBlocks = int(200/(smcStepsPerBlock))  # number of blocks until polymer configuration is saved
skipSavedBlocksBeginning = int(20/(smcStepsPerBlock))  # how many blocks (saved) to skip after you restart LEF positions
#totalSavedBlocks = 5000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
totalSavedBlocks = 4000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
#restartMilkerEveryBlocks = int(200/(smcStepsPerBlock))
restartMilkerEveryBlocks = int(400/(smcStepsPerBlock))
#Only one Hamiltonian can be loaded at a time to the simkt, but we want to update the bonds every time a LEF steps. Loading a new Hamiltonian costs a lot of time. Instead we precalculate bonds and load all positions at once as one big Hamiltonian and just change the prefactors.

# parameters for smc bonds

smcBondWiggleDist = 0.2
smcBondDist = 0.5

#if len(sys.argv)!=8:
#    print("Warning: Number of input arguments != 8")
#    sys.exit('Number of input arguments is not correct')

FLAG="" #extra label for directory name

#######use custom class to parse inputs with keywords#######################

params= tools.argsList()
for p in params.arg_dict:
    print(p, params.arg_dict[p])

if "gpu" in params.arg_dict:
    GPU = int(params.arg_dict["gpu"])
if "lifetime" in params.arg_dict:
    LIFETIME = float(params.arg_dict["lifetime"])
if "separation" in params.arg_dict:
    SEPARATION = float(params.arg_dict["separation"])
if "initiation" in params.arg_dict:
    kinPol = float(params.arg_dict["initiation"])
if "termination" in params.arg_dict:
    kterPol = float(params.arg_dict["termination"])
if "dissociation" in params.arg_dict:
    poldissoc= float(params.arg_dict["dissociation"])
if "stall" in params.arg_dict: #pol stall
    stalProb = float(params.arg_dict["stall"])
if "stallgene" in params.arg_dict:
    stall_in_gene = float(params.arg_dict["stallgene"])
if "unstall" in params.arg_dict: # pol unstall
    unstall_in_gene=float(params.arg_dict["unstall"])
if "lefspeed" in params.arg_dict:
    lef_speed=float(params.arg_dict["lefspeed"])
    pauseArray= lef_speed*pauseArray
if "lefperm" in params.arg_dict:
    lefperm=float(params.arg_dict["lefperm"])
if "shrink" in params.arg_dict:
    shrink_speed=float(params.arg_dict["shrink"])
    shrinkPauseArray= shrinkPauseArray + shrink_speed
    if shrink_speed+lef_speed>1.:
        print("WARNING: step+shrink > 1!! extrusion probabilities will not be computed correctly.\n")
if "lefstall" in params.arg_dict:
    stg = float(params.arg_dict["lefstall"])
if "lefunstall" in params.arg_dict:
    unstallLEFRate = float(params.arg_dict["lefunstall"])
if "polspeed" in params.arg_dict:
    PolSpeed = float(params.arg_dict["polspeed"])
if "polpause" in params.arg_dict:
    PolPause = float(params.arg_dict["polpause"])
if "permL" in params.arg_dict:
    L=float(params.arg_dict["permL"])
if "permR" in params.arg_dict:
    R=float(params.arg_dict["permR"])
if "collisionlife" in params.arg_dict:
    collisionLifeFactor=float(params.arg_dict["collisionlife"])
if "tssload" in params.arg_dict:
    TSSloadbias=float(params.arg_dict["tssload"])
if "tssloadstart" in params.arg_dict:
    TSSloadstart=int(params.arg_dict["tssloadstart"])
if "tssloadend" in params.arg_dict:
    TSSloadend=int(params.arg_dict["tssloadend"])
if "ttsunload" in params.arg_dict:
    TTSunload=float(params.arg_dict["ttsunload"])
if "ttszone" in params.arg_dict:
    unloadZone=int(params.arg_dict["ttszone"])
if "genelen" in params.arg_dict:
    base_genelen=int(params.arg_dict["genelen"])
if "vperm" in params.arg_dict:
    variable_permeability = int(params.arg_dict["vperm"])
if "vpermtype" in params.arg_dict:
    variable_type = int(params.arg_dict["vpermtype"])
if "vpermoffset" in params.arg_dict:
    variable_offset = float(params.arg_dict["vpermoffset"])
if "vpermpos" in params.arg_dict:
    variable_pos = int(params.arg_dict["vpermpos"])
if "vreturn" in params.arg_dict:
    variable_return = int(params.arg_dict["vreturn"])
if "vtss" in params.arg_dict:
    variable_TSSfactor = float(params.arg_dict["vtss"])
if "vgene" in params.arg_dict:
    variable_genelength = int(params.arg_dict["vgene"])
if "fixed_vgene" in params.arg_dict:
    fixed_variable_genelength=int(params.arg_dict["fixed_vgene"])
if "convergent" in params.arg_dict:
    if int(params.arg_dict["convergent"]):
        #print("warning only correct for genelength = 110 right now.\n")
        polloading=     np.array([840,  1160, 1640, 1960, 2440, 2760, 3240, 3560, 4040, 4360, 4840, 5160, 5640, 5960, 6440, 6760, 7240, 7560, 8040, 8360, 8840, 9160])
        poltermination= np.array([1000, 1001, 1800, 1801, 2600, 2601, 3400, 3401, 4200, 4201, 5000, 5001, 5800, 5801, 6600, 6601, 7400, 7401, 8200, 8201, 9000, 9001])
if "sparse" in params.arg_dict:
    if int(params.arg_dict["sparse"]):
        #print("warning only correct for genelength = 110 right now.\n")
        polloading= np.array([1950,3950,5950,7950])
        poltermination=np.array([2850,4850,6850,8850])
if "ctcfint" in params.arg_dict:
    ctcf_interval=int(params.arg_dict["ctcfint"])
if "ctcf" in params.arg_dict:
    if params.arg_dict["ctcf"] == "tss":
        #puts ctcf just before each tss, alternating left/right
        ctcf_left_list = polloading[0::2] - 1
        ctcf_right_list = polloading[1::2] - 1
    elif params.arg_dict["ctcf"] == "body":
        #puts ctcf halfway through the gene, assuming fixed gene length and genes pointing downstream
        ctcf_left_list = polloading[0::2] + base_genelen // 2
        ctcf_right_list = polloading[1::2] + base_genelen // 2
    elif params.arg_dict["ctcf"] == "distributed":
        #put in sites every ~300 kb
        Nctcf = int(N/ctcf_interval/2)
        ctcf_left_list = np.arange(1,Nctcf+2)*2*ctcf_interval-ctcf_interval
        ctcf_right_list = np.arange(1,Nctcf+1)*2*ctcf_interval
        ctcf_left_list=np.delete(ctcf_left_list, np.where(ctcf_left_list>=N)[0])
        ctcf_right_list=np.delete(ctcf_right_list, np.where(ctcf_right_list>=N)[0])
    elif params.arg_dict["ctcf"] == "distributed2":
        Nctcf = int(N/ctcf_interval/2)
        ctcf_left_list = np.arange(1,Nctcf+2)*2*ctcf_interval-int(3*ctcf_interval/2)
        ctcf_right_list = np.arange(1,Nctcf+1)*2*ctcf_interval-int(ctcf_interval/2)
        ctcf_left_list=np.delete(ctcf_left_list, np.where(ctcf_left_list>=N)[0])
        ctcf_right_list=np.delete(ctcf_right_list, np.where(ctcf_right_list>=N)[0])
    else:
        with open(params.arg_dict["ctcf"], "r") as ctcffile:
            ctcfdata=ctcffile.readlines()
            #should only be 2 lines
            entries=ctcfdata[0].split()
            ctcf_left_list = np.array([int(x) for x in entries])
            entries=ctcfdata[1].split()
            ctcf_right_list = np.array([int(x) for x in entries])
if "strongctcf" in params.arg_dict:
    strongCTCFstall= int(params.arg_dict["strongctcf"])
#granting command line control of these variables for flexibility:
if "save" in params.arg_dict:
    saveEveryBlocks = int(int(params.arg_dict["save"])/(smcStepsPerBlock))
if "skip" in params.arg_dict:
    skipSavedBlocksBeginning = int(int(params.arg_dict["skip"])/(smcStepsPerBlock))
if "total" in params.arg_dict:
    totalSavedBlocks = int(params.arg_dict["total"])
if "stallfile" in params.arg_dict:
    stallfile=params.arg_dict["stallfile"]
    STALL_FROM_FILE=True
if "restart" in params.arg_dict:
    restartMilkerEveryBlocks = int(int(params.arg_dict["restart"])/(smcStepsPerBlock))
if "log" in params.arg_dict:
    logname=params.arg_dict["log"]
if "flag" in params.arg_dict:
    FLAG=params.arg_dict["flag"]

###### a few variables to be initialized last b/c they depend on others ###
collisionLife= LIFETIME*collisionLifeFactor

pauseArrayPol=np.zeros(N,dtype=np.double) + PolSpeed ## The speed of PolII at each lattice site
for i in polloading:
    pauseArrayPol[i]=PolPause

genelength=[base_genelen]*len(polloading)
if variable_genelength:
    if fixed_variable_genelength:
        genelength=random.shuffle([80,81,82,84,86,90,95,100,110])
    else:
        for ii in range(len(genelength)):
            genelength[ii] = genelength[ii] + int(np.random.normal(0, np.sqrt(genelength[ii])))
    print("genelengths:", genelength)

#stalling, unstalling and termination in gene and after gene ends
for i,j, gl in zip(polloading,poltermination, genelength):
    if j>i:
        stalProbPol[i:i+gl]=stall_in_gene
        unstallArray[i:i+gl]=unstall_in_gene
        stalProbPol[i+gl:j]=stalProb # This line adds stall probability to all sites between end of gene (loading site + gene len) and the "termination site" (where pol stops with prob=1)
               # int(poltermination[0]-polloading[0])
        kterPolArray[i+gl:j]=kterPol
        kterPolArray[j]=kterPol
    else: # flipped genes
        stalProbPol[i-gl+1:i+1]=stall_in_gene
        unstallArray[i-gl:i+1]=unstall_in_gene
        stalProbPol[j+1:i-gl+1]=stalProb
        kterPolArray[j+1:i-gl+1]=kterPol
        kterPolArray[j]=kterPol

#permeabilities
if not variable_permeability:
        permLeftArray += L
        permRightArray += R
else:
    for ii in range(len(genelength)):
        direc = 2*int(poltermination[ii]>polloading[ii])-1
        gene_end=genelength[ii]*direc+polloading[ii]
        RL_factor=1.
        if not L==0.:
            RL_factor=R/L
        if variable_type == 0:
            for j in range(polloading[ii], gene_end):
                permLeftArray[j] = L*(gene_end-j)/(gene_end-polloading[ii])*direc + variable_offset
                permRightArray[j] = R*(gene_end-j)/(gene_end-polloading[ii])*direc + variable_offset*RL_factor
            permLeftArray[polloading[ii]] *= variable_TSSfactor
            permRightArray[polloading[ii]] *= variable_TSSfactor
            if gene_end<poltermination[ii]:
                s1=gene_end
                s2=poltermination[ii]
            else:
                s1=poltermination[ii]+1
                s2=gene_end+1
            for j in range(s1,s2):
                if not variable_return:
                    permLeftArray[j] = variable_offset
                    permRightArray[j] = variable_offset*RL_factor  # note that RL_factor adjusts right-left permeability of RNAP, not head-on/co-directional collision permeability. should change that.
                else:
                    permLeftArray[j] = L + variable_offset
                    permRightArray[j] = R + variable_offset*RL_factor
        elif variable_type == 1:
            perm_break=polloading[ii]+direc*variable_pos
            if polloading[ii] < perm_break:
                s1=polloading[ii]
                s2=perm_break
            else:
                s1=perm_break+1
                s2=polloading[ii]+1
            for j in range(s1,s2):
                permLeftArray[j] = L
                permRightArray[j] = R
            permLeftArray[polloading[ii]] *= variable_TSSfactor
            permRightArray[polloading[ii]] *= variable_TSSfactor
            if perm_break<gene_end:
                s1=perm_break
                s2=gene_end
            else:
                s1=gene_end+1
                s2=perm_break+1
            for j in range(s1,s2):
                permLeftArray[j] = variable_offset
                permRightArray[j] = variable_offset*RL_factor
            if gene_end<poltermination[ii]:
                s1=gene_end
                s2=poltermination[ii]
            else:
                s1=poltermination[ii]+1
                s2=gene_end+1
            for j in range(s1,s2):
                if not variable_return:
                    permLeftArray[j] = variable_offset
                    permRightArray[j] = variable_offset*RL_factor
                else:
                    permLeftArray[j] = L
                    permRightArray[j] = R

# -- data folder -
folder="data/"
#folder = "/net/levsha/scratch/ebanigan/txn/"

FullFileName=folder
while True:
    fname="extr_life{0}_sep{1}_density{2}_N{3}_kinpol{4}_kterPol{5}_TSSPause{6}_save{7}_total{8}_PSpeed{9}_stf{10}_stg{11}_th0.001_genelen{12}_stall{13}_L{14}_R{15}".format(LIFETIME, SEPARATION, dens,N,kinPol,kterPol,PolPause,saveEveryBlocks,totalSavedBlocks,PolSpeed,stiff,stg,base_genelen,stalProb,L,R)
    if variable_permeability:
        fname=fname+"_vperm"+str(variable_type)+"_offset"+str(variable_offset)
        if variable_type == 1:
            fname=fname+"_pos"+str(variable_pos)
        if variable_return == 1:
            fname=fname+"_vret"
        if not (variable_TSSfactor == 1.):
            fname=fname+"_vtss"+str(variable_TSSfactor)
    if not (TSSloadbias == 1.):
        fname=fname+"_tssload"+str(TSSloadbias)
        if not ((TSSloadstart==0) and (TSSloadend==0)):
            fname=fname+"_st{0}end{1}".format(TSSloadstart,TSSloadend)
    if variable_genelength:
        fname=fname+"_vgenelen"
    if not (collisionLifeFactor == 1.0):
        fname=fname+"_collife"+str(collisionLife)
    if not (TTSunload == 1.):
        fname=fname+"_ttslife"+str(TTSunload*LIFETIME)+"width"+str(unloadZone)
    if poldissoc > 0.:
        fname=fname+"_dissoc"+str(poldissoc)
    if lef_speed < 1.0:
        fname=fname+"_lefspeed"+str(lef_speed)
    if shrink_speed > 0.0:
        fname=fname+"_shrink"+str(shrink_speed)
    if stall_in_gene > 0.:
        fname=fname+"_stallgene{0}_unstall{1}".format(stall_in_gene, unstall_in_gene)
    if lefperm > 0.:
        fname=fname+"_lefperm{0}".format(lefperm)
    fname=fname+FLAG
    fname=fname+"_"+str(run_id)
    FullFileName=os.path.join(folder, fname)
    sleeptime=np.random.uniform(0,10)
    os.system("sleep {0}".format(sleeptime))
    if not os.path.exists(FullFileName):
        os.system("mkdir {0}".format(FullFileName))
        print("directory is {0}".format(FullFileName))
        break
    else:
        run_id = run_id+1

# -- Assertions to make sure parameters have been chosen correctly --

assert restartMilkerEveryBlocks % saveEveryBlocks == 0
assert (skipSavedBlocksBeginning * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert (totalSavedBlocks * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert smcStepsPerBlock<6 # max number of steps per smc block should not be too large to prevent 'jerky' polymer motion

savesPerMilker = restartMilkerEveryBlocks // saveEveryBlocks
milkerInitsSkip = saveEveryBlocks * skipSavedBlocksBeginning  // restartMilkerEveryBlocks
milkerInitsTotal  = (totalSavedBlocks + skipSavedBlocksBeginning) * saveEveryBlocks // restartMilkerEveryBlocks
print("Milker will be initialized {0} times, first {1} will be skipped".format(milkerInitsTotal, milkerInitsSkip))

# create filenames for Ekin, Epot and time

Ekin_fname = os.path.join(FullFileName,'Ekin.txt')
Epot_fname = os.path.join(FullFileName,'Epot.txt')
time_fname = os.path.join(FullFileName,'time.txt')
Par_fname  = os.path.join(FullFileName,'Pars.txt')

def save_Es_ts_Rg():
    with open(time_fname, "a+") as time_file:
        time_file.write('%f\n'%(a.state.getTime()/openmmlib.ps))
    with open(Ekin_fname, "a+") as Ekin_file:
        Ekin_file.write('%f\n'%((a.state.getKineticEnergy())/a.N/a.kT))
    with open(Epot_fname, "a+") as Epot_file:
        Epot_file.write('%f\n'%((a.state.getPotentialEnergy()) /a.N /a.kT))

class smcTranslocatorMilker(object):

    def __init__(self, smcTransObject):
        """
        :param smcTransObject: smc translocator object to work with
        """
        self.smcObject = smcTransObject
        self.allBonds = []

    def setParams(self, activeParamDict, inactiveParamDict):
        """
        A method to set parameters for bonds.
        It is a separate method because you may want to have a Simulation object already existing

        :param activeParamDict: a dict (argument:value) of addBond arguments for active bonds
        :param inactiveParamDict:  a dict (argument:value) of addBond arguments for inactive bonds

        """
        self.activeParamDict = activeParamDict
        self.inactiveParamDict = inactiveParamDict

    def setup(self, bondForce,  blocks = 100, smcStepsPerBlock = 1):
        """
        A method that milks smcTranslocator object
        and creates a set of unique bonds, etc.

        :param bondForce: a bondforce object (new after simulation restart!)
        :param blocks: number of blocks to precalculate
        :param smcStepsPerBlock: number of smcTranslocator steps per block
        :return:
        """

        if len(self.allBonds) != 0:
            raise ValueError("Not all bonds were used; {0} sets left".format(len(self.allBonds)))

        self.bondForce = bondForce

        #precalculating all bonds
        allBonds = []
        for dummy in range(blocks):
            self.smcObject.steps(smcStepsPerBlock)
            left, right = self.smcObject.getSMCs()
            bonds = [(int(i), int(j)) for i,j in zip(left, right)]
            allBonds.append(bonds)

        self.allBonds = allBonds
        self.uniqueBonds = list(set(sum(allBonds, []))) # 'sum' preserves order and makes one long list with bonds, 'set' creates a set with left bonds from different time points ordered from small to large and eliminates two equal bonds (also if they were created by different LEFs at different times). List turns set into a list with unique bonds at different time points.

        # adding forces and getting bond indices
        self.bondInds = []
        self.curBonds = allBonds.pop(0) # pop(0) removes and returns first list of bonds

        for bond in self.uniqueBonds:
            paramset = self.activeParamDict if (bond in self.curBonds) else self.inactiveParamDict
            ind = bondForce.addBond(bond[0], bond[1], **paramset)
            self.bondInds.append(ind)
        self.bondToInd = {i:j for i,j in zip(self.uniqueBonds, self.bondInds)}
        return self.curBonds,[]

    def step(self, context, verbose=False):
        """
        Update the bonds to the next step.
        It sets bonds for you automatically!
        :param context:  context
        :return: (current bonds, previous step bonds); just for reference
        """
        if len(self.allBonds) == 0:
            raise ValueError("No bonds left to run; you should restart simulation and run setup  again")

        pastBonds = self.curBonds
        self.curBonds = self.allBonds.pop(0)  # getting current bonds
        bondsRemove = [i for i in pastBonds if i not in self.curBonds]
        bondsAdd = [i for i in self.curBonds if i not in pastBonds]
        bondsStay = [i for i in pastBonds if i in self.curBonds]
        if verbose:
            print("{0} bonds stay, {1} new bonds, {2} bonds removed".format(len(bondsStay),
                                                                            len(bondsAdd), len(bondsRemove)))
        bondsToChange = bondsAdd + bondsRemove
        bondsIsAdd = [True] * len(bondsAdd) + [False] * len(bondsRemove)
        for bond, isAdd in zip(bondsToChange, bondsIsAdd):
            ind = self.bondToInd[bond]
            paramset = self.activeParamDict if isAdd else self.inactiveParamDict
            self.bondForce.setBondParameters(ind, bond[0], bond[1], **paramset)  # actually updating bonds
        self.bondForce.updateParametersInContext(context)  # now run this to update things in the context
        return self.curBonds, pastBonds

    def getRNAP(self):
        return self.smcObject.getPolOccupied()

def initModel():
    birthArray = np.zeros(N,dtype=np.double) + 0.1
    #birthArray[polloading] *= TSSloadbias
    for nn, tt  in zip(polloading,poltermination):
        #the following accounts for different gene directions:
        if tt > nn:
            loadbegin=min(max(nn-TSSloadstart,0), N) #max/min construction here ensures range will not exceed ends of polymer, noting that TSSloadstart can be + or -
            loadend=max(min(nn+TSSloadend+1,N),0)
        else:
            loadbegin= min(max(nn-TSSloadend,0),N)
            loadend= max(min(nn+TSSloadstart+1,N),0)
        for jj in range(loadbegin,loadend):
            birthArray[jj] *= TSSloadbias
    deathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME)
    collisionDeathArray= np.zeros(N,dtype=np.double) + 1./collisionLife
    if not (TTSunload==1.):
        for nn,tt,gg in zip(polloading, poltermination, genelength):
            if tt>nn:
                zonebegin=nn+gg
                zoneend=min(nn+gg+unloadZone,N)
            else:
                zonebegin= max(nn-gg+1-unloadZone,0)
                zoneend= nn-gg+1
            for jj in range(zonebegin, zoneend):
                deathArray[jj] = 1./(TTSunload*LIFETIME)
                collisionDeathArray[jj] =  1./(TTSunload*LIFETIME)

    stallLeftArray = np.zeros(N, dtype=np.double) #stall prob for left LEF legs, i.e. CTCF sites that point rightward
    stallRightArray = np.zeros(N, dtype=np.double)
    if STALL_FROM_FILE: #read in stall rates from file; file should provide list of LEF stalling rates starting from the TSS, extending as long as the configuration of repeats permits
        with open("stallfiles/"+stallfile,"r") as infile:
            stall_list = infile.readlines()
        for entry in stall_list:
            gene_stall.append(float(entry.split()[0]))
        for nn,tt in zip(polloading,poltermination):
            ii=0
            if nn<tt:
                gene_direc=1
            else:
                gene_direc=-1
            for ii in range(len(gene_stall)):
                stallLeftArray[nn+ii*gene_direc]=gene_stall[ii]
                stallRightArray[nn+ii*gene_direc]=gene_stall[ii]
    else: #stall rates set by stall params & CTCF position options
        stallLeftArray[ctcf_left_list] = stg
        stallRightArray[ctcf_right_list] = stg

    unstallLEFArray = np.ones(N, dtype=np.double) * unstallLEFRate # generally have a single unstall rate.  doesn't matter if non-stall sites have unstall>0, since there won't be stalls there anyway
    stallDeathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME) # unbinding rate during stall can be different, but we usually leave it the same
    smcNum = int(N / SEPARATION)
    curPos = 0

    SMCTran = smcTranslocatorDirectional(birthArray, deathArray, stallLeftArray, stallRightArray, unstallLEFArray, pauseArray, stallDeathArray, smcNum,
                                         kinPol,kterPolArray,pauseArrayPol,shrinkPauseArray, polloading,poltermination, stalProbPol, unstallArray,
                                         PolPermL=permLeftArray, PolPermR=permRightArray, collisionFalloffProb=collisionDeathArray,
                                         poldissoc= poldissoc,
                                         LefPerm=lefperm,
                                         strongCTCF=strongCTCFstall
                                        )
    return SMCTran

SMCTran = initModel()  # defining actual smc translocator object

# now polymer simulation code starts

# ------------feed smcTran to the milker---
SMCTran.steps(1000000)  # first steps to "equilibrate" SMC dynamics. If desired of course.
milker = smcTranslocatorMilker(SMCTran)   # now feed this thing to milker (do it once!)
#--------- end new code ------------

for milkerCount in range(milkerInitsTotal):
    doSave = milkerCount >= milkerInitsSkip

    # simulation parameters are defined below
    a = Simulation(timestep=80, thermostat=0.001)#Collision rate in inverse picoseconds, low collistion rate means ballistic like motion, default in openmmpolymer is 0.001. Motion polymer is not diffusive, this is ok for statistical average,
    #but not for dynamics of the polymer
    a.setup(platform="CUDA", PBC=True, PBCbox=[box, box, box], GPU=GPU, precision="mixed")  # set up GPU here, PBC=Periodic Boundary Conditions. Default integrator is langevin with 300 K, friction coefficient of 1/ps, step size 0.002ps
    a.saveFolder(FullFileName)
    a.load(data)
    a.addHarmonicPolymerBonds(wiggleDist=0.1) # WiggleDist controls distance at which energy of bond equals kT
    if stiff > 0:
        a.addGrosbergStiffness(stiff) # Chain stiffness is introduced by an angular potential U(theta)=stiff(1-cos(theta-Pi))
    a.addPolynomialRepulsiveForce(trunc=1.5, radiusMult=1.05) #Polynomial repulsive potential between particles. Has value trunc=3.0 at zero, stays flat until 0.6-0.7 and then drops to zero. For attraction between a selective set of particles, use LeonardJones or addSelectiveSSWForce (see blocks.py or ask Johannes)
    a.step = block

    # ------------ initializing milker; adding bonds ---------
    # copied from addBond
    kbond = a.kbondScalingFactor / (smcBondWiggleDist ** 2)
    bondDist = smcBondDist * a.length_scale

    activeParams = {"length":bondDist,"k":kbond}
    inactiveParams = {"length":bondDist, "k":0}
    milker.setParams(activeParams, inactiveParams)

    # this step actually puts all bonds in and sets first bonds to be what they should be
    milker.setup(bondForce=a.forceDict["HarmonicBondForce"],
                 blocks=restartMilkerEveryBlocks,   # default value; milk for 100 blocks
                 smcStepsPerBlock=smcStepsPerBlock)  #
    print("Restarting milker")

    a.doBlock(steps=steps, increment=False)  # do block for the first time with first set of bonds in
    #print('done 1')
    for i in range(restartMilkerEveryBlocks - 1):
        #print(i)
        curBonds, pastBonds = milker.step(a.context)  # this updates bonds. You can do something with bonds here
        if i % saveEveryBlocks == (saveEveryBlocks - 2):
            a.doBlock(steps=steps, increment = doSave)
            if doSave:
                a.save()
                pickle.dump(curBonds, open(os.path.join(a.folder, "SMC{0}.dat".format(a.step)),'wb'))
                save_Es_ts_Rg() # save energies and time
                if saveRNAP:
                    pickle.dump(milker.getRNAP(), open(os.path.join(a.folder, "RNAP{0}.dat".format(a.step)), 'wb'))
        else:
            a.integrator.step(steps)  # do steps without getting the positions from the GPU (faster)

    data = a.getData()  # save data and step, and delete the simulation
    block = a.step
    del a

    time.sleep(0.2)  # wait 200ms for sanity (to let garbage collector do its magic)

with open(Par_fname,"a+") as Parfile:
    Parfile.write(" tau="+str(LIFETIME)+"\n Separation="+str(SEPARATION)+"\n N="+str(N)+"\n smcStepsPerBlock="+str(smcStepsPerBlock)+"\n stiff="+str(stiff)+"\n dens="+str(dens)+"\n block="+str(block)+"\n  SaveEveryBlocks="+str(saveEveryBlocks)+"\n skipSavedBlocksBeginning="+str(skipSavedBlocksBeginning)+"\n totalSavedBlocks="+str(totalSavedBlocks)+"\n restartMilkerEveryBlocks="+str(restartMilkerEveryBlocks)+"\n smcBondWiggleDist="+str(smcBondWiggleDist)+"\n smcBondDist="+str(smcBondDist)+"\n SmcTimestep=1\n NumMonos"+str(N))

os.system("mv {0} {1}".format(logname, FullFileName))


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py,md
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.13.3
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
# Uncomment to run the notebook in Colab
# # ! pip install -q "wax-ml[complete]@git+https://github.com/eserie/wax-ml.git"
# # ! pip install -q --upgrade jax jaxlib==0.1.70+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
# -

# %matplotlib inline

import io
import warnings
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable, NamedTuple, Optional, TypeVar

import haiku as hk
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpy as onp
import optax
import pandas as pd
import plotnine as gg
import requests
from sklearn.preprocessing import MinMaxScaler
from tqdm.auto import tqdm

from wax.accessors import register_wax_accessors
from wax.compile import jit_init_apply
from wax.encode import Encoder
from wax.modules import Buffer, FillNanInf, Lag, RollingMean
from wax.unroll import unroll

print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
jax.devices()

# # 🔭 Reconstructing the light curve of stars with LSTM 🔭
#
# [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eserie/wax-ml/blob/main/docs/notebooks/05_reconstructing_the_light_curve_of_stars.ipynb)

# Let's take a walk through the stars...
#
# This notebook is based on the study done in
# [this post by Christophe Pere](https://towardsdatascience.com/how-to-use-deep-learning-for-time-series-forecasting-3f8a399cf205)
# and the notebook available on
# [the authors's github](https://github.com/Christophe-pere/Time_series_RNN).
#
# We will repeat this study on starlight using the LSTM architecture to predict the observed light flux through time.
#
# Our LSTM implementation is based on this [notebook from Haiku's github repository](https://github.com/deepmind/dm-haiku/blob/master/examples/haiku_lstms.ipynb).
#
# We'll see how to use WAX-ML to ease the preparation of time series data stored in dataframes and having Nans
# before calling a "standard" deep-learning workflow.
#
# ## Disclaimer
#
# Despite the fact that this code works with real data, the results presented here should not be considered as scientific knowledge insights, to the knowledge of the authors of WAX-ML, neither the results nor the data source have been reviewed by an astrophysics pair.
#
# The purpose of this notebook is only to demonstrate how WAX-ML can be used when applying a "standard" machine learning workflow, here LSTM, to analyze time series.

# ## Download the data

register_wax_accessors()

# + tags=["parameters"]
# Parameters
STAR = "007609553"
SEQ_LEN = 64
BATCH_SIZE = 8
TRAIN_SIZE = 2 ** 16
NUM_EPOCHS = 10
NUM_STARS = None
RECORD_FREQ = 100
TOTAL_LEN = None
TRAIN_DATE = "2016"
CACHE_DIR = Path("./cached_data/")
# -

# %%time
filename = CACHE_DIR / "kep_lightcurves.parquet"
try:
    raw_dataframe = pd.read_parquet(open(filename, "rb"))
    print(f"data read from {filename}")
except FileNotFoundError:
    # Downloading the csv file from Chrustioge Pere GitHub account
    download = requests.get(
        "https://raw.github.com/Christophe-pere/Time_series_RNN/master/kep_lightcurves.csv"
    ).content
    raw_dataframe = pd.read_csv(io.StringIO(download.decode("utf-8")))
    # set date index
    raw_dataframe.index = pd.Index(
        pd.date_range("2009-03-07", periods=len(raw_dataframe.index), freq="h"),
        name="time",
    )
    # save dataframe locally in CACHE_DIR
    CACHE_DIR.mkdir(exist_ok=True)
    raw_dataframe.to_parquet(filename)
    print(f"data saved in {filename}")

# shortening of data to speed up the execution of the notebook in the CI
if TOTAL_LEN:
    raw_dataframe = raw_dataframe.iloc[:TOTAL_LEN]

# Let's visualize the description of this dataset:

raw_dataframe.describe().T.to_xarray()

stars = raw_dataframe.columns
stars = sorted(list(set([i.split("_")[0] for i in stars])))
print(f"The number of stars available is: {len(stars)}")
print(f"star identifiers: {stars}")

dataframe = raw_dataframe[[i + "_rscl" for i in stars]].rename(
    columns=lambda c: c.replace("_rscl", "")
)
dataframe.columns.names = ["star"]
dataframe.shape

if NUM_STARS:
    columns = dataframe.columns.tolist()
    columns.remove(STAR)
    dataframe = dataframe[[STAR] + columns[: NUM_STARS - 1]]

# ## Rolling mean

# We will smooth the data by applying a rolling mean with a window of 100 periods.

# ### Count nan values
#
# But before since the dataset has some nan values, we will extract few statistics
# about the density of nan values in windows of size 100.
#
# It will be the occasion to show a usage of the `wax.modules.Buffer` module with the `format_outputs=False`
# option for the dataframe accessor `.wax.stream`.

# Let's apply the `Buffer` module to the data:

buffer, _ = dataframe.wax.stream(format_outputs=False).apply(lambda x: Buffer(100)(x))

assert isinstance(buffer, jnp.ndarray)

# Equivalently, we can use wax `unroll` function.

buffer = unroll(lambda x: Buffer(100)(x))(jax.device_put(dataframe.values))

# Let's describe the statistic of nans with pandas:

count_nan = jnp.isnan(buffer).sum(axis=1)
pd.DataFrame(onp.array(count_nan)).stack().describe().astype(int)

# ### Computing the rolling mean

# We will choose a `min_periods` of 5 in order to keep at leas 75% of the points.

# %%time
dataframe_mean, _ = dataframe.wax.stream().apply(
    lambda x: RollingMean(100, min_periods=5)(x)
)

dataframe.iloc[:, :2].plot()

# ## Forecasting with Machine Learning
#
# We need two forecast in this data, if you look with attention you'll see micro holes and big holes.

T = TypeVar("T")

class Pair(NamedTuple):
    x: T
    y: T

class TrainSplit(NamedTuple):
    train: T
    validation: T

gg.theme_set(gg.theme_bw())
warnings.filterwarnings("ignore")

plt.rcParams["figure.figsize"] = 18, 8
fig, (ax, lax) = plt.subplots(ncols=2, gridspec_kw={"width_ratios": [4, 1]})
dataframe.plot(ax=ax, title="raw data")
ax.legend(bbox_to_anchor=(0, 0, 1, 1), bbox_transform=lax.transAxes)
lax.axis("off")

plt.rcParams["figure.figsize"] = 18, 8
fig, (ax, lax) = plt.subplots(ncols=2, gridspec_kw={"width_ratios": [4, 1]})
dataframe_mean.plot(ax=ax, title="Smoothed data")
ax.legend(bbox_to_anchor=(0, 0, 1, 1), bbox_transform=lax.transAxes)
lax.axis("off")
# -

# ### Normalize data

dataframe_mean.stack().hist(bins=100, log=True)

def min_max_scaler(values: pd.DataFrame, output_format: str = "dataframe") -> Encoder:
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaler.fit(values)
    index = values.index
    columns = values.columns

    def encode(dataframe: pd.DataFrame):
        nonlocal index
        nonlocal columns

        index = dataframe.index
        columns = dataframe.columns
        array_normed = scaler.transform(dataframe)

        if output_format == "dataframe":
            return pd.DataFrame(array_normed, index, columns)
        elif output_format == "jax":
            return jnp.array(array_normed)
        else:
            return array_normed

    def decode(array_scaled):

        value = scaler.inverse_transform(array_scaled)

        if output_format == "dataframe":
            return pd.DataFrame(value, index, columns)
        else:
            return value

    return Encoder(encode, decode)

# -

scaler = min_max_scaler(dataframe_mean)
dataframe_normed = scaler.encode(dataframe_mean)
assert (scaler.decode(dataframe_normed) - dataframe_mean).stack().abs().max() < 1.0e-4

dataframe_normed.stack().hist(bins=100)

# ### Prepare train / validation datasets

def split_feature_target(
    dataframe,
    look_back=SEQ_LEN,
    shuffle=True,
    stack=True,
    min_periods_ratio: float = 0.8,
    rng=None,
) -> Pair:
    def prepare_xy(data):
        buffer = Buffer(look_back + 1)(data)
        x = buffer[:-1]
        y = buffer[-1]
        return x, y

    def prepare_xy(data):
        y = Buffer(look_back)(data)
        x = Lag(1)(y)
        return x, y

    x, y = unroll(prepare_xy)(jax.device_put(dataframe.values))

    if shuffle:
        if rng is None:
            rng = jax.random.PRNGKey(42)

        B = x.shape[0]
        idx = jnp.arange(B)
        idx = jax.random.shuffle(rng, idx)

        x = x[idx]
        y = y[idx]

    if stack:
        B, T, F = x.shape
        x = x.transpose(1, 0, 2).reshape(T, B * F, 1).transpose(1, 0, 2)
        y = y.transpose(1, 0, 2).reshape(T, B * F, 1).transpose(1, 0, 2)

    if min_periods_ratio:
        T = x.shape[1]
        count_nan = jnp.isnan(x).sum(axis=1)
        mask = count_nan < min_periods_ratio * T
        idx = jnp.where(mask)
        x = x[idx[0]]
        y = y[idx[0]]

    # round Batch size to a power of to
    B = x.shape[0]
    B_round = int(2 ** jnp.floor(jnp.log2(B)))
    print(f"{B} batches rounded to {B_round} batches.")
    x = x[:B_round]
    y = y[:B_round]

    # fillnan by zeros
    x, y = hk.testing.transform_and_run(lambda x: FillNanInf()(x))((x, y))

    return Pair(x, y)

# +
# split_feature_target(dataframe)
# -

def split_train_validation(
    dataframe, train_size, look_back, scaler: Optional[Callable] = None
) -> TrainSplit:

    # prepare scaler
    train_df = dataframe.iloc[:train_size]

    if scaler:
        scaler = scaler(train_df)

    # prepare train data
    if scaler:
        train_df = scaler.encode(train_df)

    train_xy = split_feature_target(train_df, look_back)

    # prepare validation data
    valid_size = len(dataframe) - train_size
    valid_size = int(2 ** jnp.floor(jnp.log2(valid_size)))

    valid_end = int(train_size + valid_size)
    valid_df = dataframe.iloc[train_size:valid_end]

    if scaler:
        valid_df = scaler.encode(valid_df)

    valid_xy = split_feature_target(valid_df, look_back)

    return TrainSplit(train_xy, valid_xy)

TRAIN_SIZE

print(f"Look at star: {STAR}")
train, valid = split_train_validation(dataframe_normed[[STAR]], TRAIN_SIZE, SEQ_LEN)

train[0].shape, train[1].shape, valid[0].shape, valid[1].shape

# TRAIN_SIZE, VALID_SIZE = len(train.x), len(valid.x)
print(
    f"effective train_size = {len(train.x)}, " f"effective valid size= {len(valid.x)}"
)

# Plot an observation/target pair.
rng = jax.random.PRNGKey(42)
batch_plot = jax.random.choice(rng, len(train[0]))
df = pd.DataFrame(
    {"x": train.x[batch_plot, :, 0], "y": train.y[batch_plot, :, 0]}
).reset_index()
df = pd.melt(df, id_vars=["index"], value_vars=["x", "y"])
plot = (
    gg.ggplot(df)
    + gg.aes(x="index", y="value", color="variable")
    + gg.geom_line()
    + gg.scales.scale_y_log10()
)
_ = plot.draw()

# ### Dataset iterator

class Dataset:
    """An iterator over a numpy array, revealing batch_size elements at a time."""

    def __init__(self, xy: Pair, batch_size: int):
        self._x, self._y = xy
        self._batch_size = batch_size
        self._length = self._x.shape[0]
        self._idx = 0
        if self._length % batch_size != 0:
            msg = "dataset size {} must be divisible by batch_size {}."
            raise ValueError(msg.format(self._length, batch_size))

    def __next__(self) -> Pair:
        start = self._idx
        end = start + self._batch_size
        x, y = self._x[start:end], self._y[start:end]
        if end >= self._length:
            print(f"End of the data set (size={end}). Return to the beginning.")
            end = end % self._length
            assert end == 0  # Guaranteed by ctor assertion.
        self._idx = end
        return Pair(x, y)

# + [markdown] colab_type="text" id="LZGw5Jdvjmqh"
# ### Training an LSTM
#
# To train the LSTM, we define a Haiku function which unrolls the LSTM over the input sequence, generating predictions for all output values. The LSTM always starts with its initial state at the start of the sequence.
#
# The Haiku function is then transformed into a pure function through `hk.transform`, and is trained with Adam on an L2 prediction loss.

# + colab={} colab_type="code" id="nacnTj5ejIK5"
def unroll_net(seqs: jnp.ndarray):
    """Unrolls an LSTM over seqs, mapping each output to a scalar."""
    # seqs is [T, B, F].
    core = hk.LSTM(32)
    batch_size = seqs.shape[0]
    outs, state = hk.dynamic_unroll(
        core, seqs, core.initial_state(batch_size), time_major=False
    )
    # We could include this Linear as part of the recurrent core!
    # However, it's more efficient on modern accelerators to run the linear once
    # over the entire sequence than once per sequence element.
    return hk.BatchApply(hk.Linear(1))(outs), state

# + colab={} colab_type="code" id="nacnTj5ejIK5"
model = jit_init_apply(hk.transform(unroll_net))

# +
@jax.jit
def loss(pred, y):
    return jnp.mean(jnp.square(pred - y))

def model_with_loss(x, y):
    pred, _ = unroll_net(x)
    return loss(pred, y)

# + colab={} colab_type="code" id="nacnTj5ejIK5"
class TrainState(NamedTuple):
    step: int
    params: Any
    opt_state: Any
    rng: jnp.ndarray
    loss: float

def train_model(
    model_with_loss: Callable,
    train_ds: Dataset,
    valid_ds: Dataset,
    max_iterations: int = -1,
    rng=None,
    record_freq=100,
) -> hk.Params:
    """Initializes and trains a model on train_ds, returning the final params."""
    opt = optax.adam(1e-3)
    model_with_loss = jit_init_apply(hk.transform(model_with_loss))

    @jax.jit
    def update(train_state, x, y):
        step, params, opt_state, rng, _ = train_state
        if rng is not None:
            (rng,) = jax.random.split(rng, 1)
        l, grads = jax.value_and_grad(model_with_loss.apply)(params, rng, x, y)
        grads, opt_state = opt.update(grads, opt_state)
        params = optax.apply_updates(params, grads)
        return TrainState(step + 1, params, opt_state, rng, l)

    # Initialize state.
    def init():
        x, y = next(train_ds)
        params = model_with_loss.init(rng, x, y)
        opt_state = opt.init(params)
        return TrainState(0, params, opt_state, rng, jnp.inf)

    def _format_results(records):
        records = {key: jnp.stack(l) for key, l in records.items()}
        return records

    records = defaultdict(list)
    train_state = init()
    with tqdm(total=max_iterations if max_iterations > 0 else None) as pbar:
        while True:
            try:
                x, y = next(train_ds)
            except StopIteration:
                return train_state, _format_results(records)

            train_state = update(train_state, x, y)
            if train_state.step % record_freq == 0:
                x, y = next(valid_ds)
                if rng is not None:
                    (rng,) = jax.random.split(rng, 1)
                valid_loss = model_with_loss.apply(train_state.params, rng, x, y)
                records["step"].append(train_state.step)
                records["valid_loss"].append(valid_loss)
                records["train_loss"].append(train_state.loss)

            pbar.update()
            if max_iterations > 0 and train_state.step >= max_iterations:
                return train_state, _format_results(records)

# + colab={} colab_type="code" id="AssgDctokbl5"
# %%time
train, valid = split_train_validation(dataframe_normed[[STAR]], TRAIN_SIZE, SEQ_LEN)
train_ds = Dataset(train, BATCH_SIZE)
valid_ds = Dataset(valid, BATCH_SIZE)

train_state, records = train_model(
    model_with_loss,
    train_ds,
    valid_ds,
    len(train.x) // BATCH_SIZE * NUM_EPOCHS,
    rng=jax.random.PRNGKey(42),
    record_freq=RECORD_FREQ,
)

# +
# train_state.params
# -

# Plot losses
losses = pd.DataFrame(records)
df = pd.melt(losses, id_vars=["step"], value_vars=["train_loss", "valid_loss"])
plot = (
    gg.ggplot(df)
    + gg.aes(x="step", y="value", color="variable")
    + gg.geom_line()
    + gg.scales.scale_y_log10()
)
_ = plot.draw()

# + [markdown] colab_type="text" id="yr7jrOL3ki-b"
# ### Sampling
#
# The point of training models is so that they can make predictions! How can we generate predictions with the trained model?
#
# If we're allowed to feed in the ground truth, we can just run the original model's `apply` function.

# + colab={} colab_type="code" id="f2qETEqXLT1N"
def plot_samples(truth: np.ndarray, prediction: np.ndarray) -> gg.ggplot:
    assert truth.shape == prediction.shape
    df = pd.DataFrame(
        {"truth": truth.squeeze(), "predicted": prediction.squeeze()}
    ).reset_index()
    df = pd.melt(df, id_vars=["index"], value_vars=["truth", "predicted"])
    plot = (
        gg.ggplot(df) + gg.aes(x="index", y="value", color="variable") + gg.geom_line()
    )
    return plot

# + colab={} colab_type="code" id="KOuK1egilGD0"
# Grab a sample from the validation set.
sample_x, sample_y = next(valid_ds)
sample_x = sample_x[:1]  # Shrink to batch-size 1.
sample_y = sample_y[:1]

# Generate a prediction, feeding in ground truth at each point as input.
predicted, _ = model.apply(train_state.params, None, sample_x)

plot = plot_samples(sample_y, predicted)
plot.draw()
del sample_x, predicted

# -

# ### Run autoregressively

# + [markdown] colab_type="text" id="tDyGshz_lwrM"
# If we can't feed in the ground truth (because we don't have it), we can also run the model autoregressively.

# + colab={} colab_type="code" id="Cg8oQ75Ulvld"
def autoregressive_predict(
    trained_params: hk.Params,
    context: jnp.ndarray,
    seq_len: int,
    pbar=False,
):
    """Given a context, autoregressively generate the rest of a sine wave."""

    ar_outs = []
    context = jax.device_put(context)
    times = onp.arange(seq_len - context.shape[1] + 1)
    if pbar:
        times = tqdm(times)
    for _ in times:
        full_context = jnp.concatenate([context] + ar_outs, axis=1)

        outs, _ = model.apply(trained_params, None, full_context)
        # Append the newest prediction to ar_outs.
        ar_outs.append(outs[:, -1:, :])
    # Return the final full prediction.
    return outs

# + colab={} colab_type="code" id="Cg8oQ75Ulvld"
sample_x, sample_y = next(valid_ds)
sample_x = sample_x[:1]  # Shrink to batch-size 1.
sample_y = sample_y[:1]  # Shrink to batch-size 1.

context_length = SEQ_LEN // 8
print(f"context_length = {context_length}")
# Cut the batch-size 1 context from the start of the sequence.
context = sample_x[:, :context_length]

# + colab={} colab_type="code" id="Cg8oQ75Ulvld"
# %%time
# We can reuse params we got from training for inference - as long as the
# declaration order is the same.
predicted = autoregressive_predict(train_state.params, context, SEQ_LEN, pbar=True)
# -

sample_y.shape, predicted.shape

# + colab={} colab_type="code" id="Cg8oQ75Ulvld"
plot = plot_samples(sample_y, predicted)
plot += gg.geom_vline(xintercept=context.shape[1], linetype="dashed")
_ = plot.draw()

# + [markdown] colab_type="text" id="qGkr2gf2oALo"
# #### Sharing parameters with a different function.
#
# Unfortunately, this is a bit slow - we're doing O(N^2) computation for a sequence of length N.
#
# It'd be better if we could do the autoregressive sampling all at once - but we need to write a new Haiku function for that.
#
# We're in luck - if the Haiku module names match, the same parameters can be used for multiple Haiku functions.
#
# This can be achieved through a combination of two techniques:
#
# 1. If we manually give a unique name to a module, we can ensure that the parameters are directed to the right places.
# 2. If modules are instantiated in the same order, they'll have the same names in different functions.
#
# Here, we rely on method #2 to create a fast autoregressive prediction.

# + colab_type="text" id="qGkr2gf2oALo"
@hk.transform
def fast_autoregressive_predict_fn(context, seq_len):
    """Given a context, autoregressively generate the rest of a sine wave."""
    core = hk.LSTM(32)
    dense = hk.Linear(1)
    state = core.initial_state(context.shape[0])
    # Unroll over the context using `hk.dynamic_unroll`.
    # As before, we `hk.BatchApply` the Linear for efficiency.
    context_outs, state = hk.dynamic_unroll(
        core,
        context,
        state,
        time_major=False,
    )
    context_outs = hk.BatchApply(dense)(context_outs)

    # Now, unroll one step at a time using the running recurrent state.
    ar_outs = []
    x = context_outs[:, -1, :]
    times = range(seq_len - context.shape[1])
    for _ in times:
        x, state = core(x, state)
        x = dense(x)
        ar_outs.append(x)
    ar_outs = jnp.stack(ar_outs)
    ar_outs = ar_outs.transpose(1, 0, 2)
    return jnp.concatenate([context_outs, ar_outs], axis=1)

fast_autoregressive_predict = jax.jit(
    fast_autoregressive_predict_fn.apply, static_argnums=(3,)
)

# + colab={} colab_type="code" id="WdKcHr6_n_ba"
# %%time
# Reuse the same context from the previous cell.
predicted = fast_autoregressive_predict(train_state.params, None, context, SEQ_LEN)

# + colab={} colab_type="code" id="WdKcHr6_n_ba"
# The plots should be equivalent!
plot = plot_samples(sample_y, predicted)
plot += gg.geom_vline(xintercept=context.shape[1], linetype="dashed")
_ = plot.draw()
# -

# # Sample trajectories

# + colab={} colab_type="code" id="WdKcHr6_n_ba"
sample_x, sample_y = next(valid_ds)
sample_x = sample_x[:1]  # Shrink to batch-size 1.
sample_y = sample_y[:1]  # Shrink to batch-size 1.

context_length = SEQ_LEN // 8
print(f"context_length = {context_length}")
# Cut the batch-size 1 context from the start of the sequence.
context = sample_x[:, :context_length]

# Reuse the same context from the previous cell.
predicted = fast_autoregressive_predict(train_state.params, None, context, SEQ_LEN)

# The plots should be equivalent!
plot = plot_samples(sample_y, predicted)
plot += gg.geom_vline(xintercept=context.shape[1], linetype="dashed")
_ = plot.draw()
# -

# ## timeit

# + colab={} colab_type="code" id="9S0tkPXGrU3a"
# %timeit autoregressive_predict(train_state.params, context, SEQ_LEN)
# %timeit fast_autoregressive_predict(train_state.params, None, context, SEQ_LEN)
# -
# ## Train all stars

# ### Training

def split_train_validation_date(dataframe, date, look_back) -> TrainSplit:
    train_size = len(dataframe.loc[:date])
    return split_train_validation(dataframe, train_size, look_back)

# %%time
train, valid = split_train_validation_date(dataframe_normed, TRAIN_DATE, SEQ_LEN)
print(f"effective train size = {train[0].shape[1]}")

train[0].shape, train[1].shape, valid[0].shape, valid[1].shape

train_ds = Dataset(train, BATCH_SIZE)
valid_ds = Dataset(valid, BATCH_SIZE)
# del train, valid  # Don't leak temporaries.

# %%time
train_state, records = train_model(
    model_with_loss,
    train_ds,
    valid_ds,
    len(train.x) // BATCH_SIZE * 1,
    jax.random.PRNGKey(42),
    record_freq=RECORD_FREQ,
)

# Plot losses
losses = pd.DataFrame(records)
df = pd.melt(losses, id_vars=["step"], value_vars=["train_loss", "valid_loss"])
plot = (
    gg.ggplot(df)
    + gg.aes(x="step", y="value", color="variable")
    + gg.geom_line()
    + gg.scales.scale_y_log10()
)
_ = plot.draw()

# ### Sampling

# +
# Grab a sample from the validation set.
sample_x, sample_y = next(valid_ds)
sample_x = sample_x[:1]  # Shrink to batch-size 1.
sample_y = sample_y[:1]  # Shrink to batch-size 1.

# Generate a prediction, feeding in ground truth at each point as input.
predicted, _ = model.apply(train_state.params, None, sample_x)

plot = plot_samples(sample_y, predicted)
_ = plot.draw()
# -

# ### Run autoregressively

# +
# %%time
sample_x, sample_y = next(valid_ds)
sample_x = sample_x[:1]  # Shrink to batch-size 1.
sample_y = sample_y[:1]  # Shrink to batch-size 1.

context_length = SEQ_LEN // 8
# Cut the batch-size 1 context from the start of the sequence.
context = sample_x[:, :context_length]

# Reuse the same context from the previous cell.
predicted = fast_autoregressive_predict(train_state.params, None, context, SEQ_LEN)

# The plots should be equivalent!
plot = plot_samples(sample_y, predicted)
plot += gg.geom_vline(xintercept=len(context), linetype="dashed")
_ = plot.draw()


import os
from sympy import *
import pandas as pd
import numpy as np
import scipy.fftpack
import matplotlib as mpl
import matplotlib.pyplot as plt
plt.style.use("seaborn-paper")
from mpl_toolkits.axes_grid1 import make_axes_locatable

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx

def find_nearest_value(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return array[idx]

# ## Canvas palette

# In[2]:

#Canvas for single plot
x = np.linspace(0,10,100)
y = np.sin(x)
plt.figure(figsize=[14,6])
plt.grid(True)
plt.title("Change-me!",fontsize=20)
plt.plot(x,y,label="testvalue")
plt.legend(fontsize=16)
plt.xlabel("XLABEL (unit)",fontsize=18)
plt.ylabel("YLABEL (unit)",fontsize=18)
plt.show()

# In[3]:

#Canvas for side by side
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(14,6))
fig.suptitle("test",y=1.05,fontsize=20)

axes[0].grid(True)
axes[0].plot(x,y,label="testvalue")
axes[0].legend(fontsize=16)
axes[0].set_title("TESTTITLE",fontsize=18)
axes[0].set_xlabel("XLABEL (unit)",fontsize=18)
axes[0].set_ylabel("YLABEL (unit)",fontsize=18)
axes[0].legend(fontsize=16)
axes[0].tick_params(axis='both', which='major', labelsize=15)

axes[1].grid(True)
axes[1].plot(x,y,label="testvalue")
axes[1].legend(fontsize=16)
axes[1].set_title("TESTTITLE",fontsize=18)
axes[1].set_xlabel("XLABEL (unit)",fontsize=18)
axes[1].set_ylabel("YLABEL (unit)",fontsize=18)
axes[1].legend(fontsize=16)
axes[1].tick_params(axis='both', which='major', labelsize=15)

fig.tight_layout()
plt.show()

# In[4]:

#Canvas for side by side
fig, axes = plt.subplots(nrows=2, ncols=4, figsize=(14,6))
fig.suptitle("test",y=1.05,fontsize=20)

axes[0,0].grid(True)
axes[0,0].plot(x,y,label="testvalue")
axes[0,0].legend(fontsize=16)
axes[0,0].set_title("TESTTITLE",fontsize=18)
axes[0,0].set_xlabel("XLABEL (unit)",fontsize=18)
axes[0,0].set_ylabel("YLABEL (unit)",fontsize=18)
axes[0,0].legend(fontsize=16)
axes[0,0].tick_params(axis='both', which='major', labelsize=15)

axes[0,1].grid(True)
axes[0,1].plot(x,y,label="testvalue")
axes[0,1].legend(fontsize=16)
axes[0,1].set_title("TESTTITLE",fontsize=18)
axes[0,1].set_xlabel("XLABEL (unit)",fontsize=18)
axes[0,1].set_ylabel("YLABEL (unit)",fontsize=18)
axes[0,1].legend(fontsize=16)
axes[0,1].tick_params(axis='both', which='major', labelsize=15)

fig.tight_layout()
plt.show()

# ## Read data

# In[2]:

#Folder and paths definitions
main_path  = os.getcwd()
datafolder_path = main_path+"/results"
results_dir = "/output_py"
output_dir = main_path+results_dir
try:
    os.mkdir(output_dir)
except OSError:
    print ("Creation of the directory %s failed" % results_dir)
else:
    print ("Successfully created the directory %s " % results_dir)

# In[27]:

Kvalues = np.linspace(0,6,30)
Kvalues = np.around(Kvalues, decimals=3)

# In[28]:

pvalues = [0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, .95, 1]

# In[9]:

#Create dataframe dictionary. For each entry, first value is the K of the dataframe (second value)
data = []
for i in range(0,len(Kvalues)):
    for j in pvalues:
        filename = datafolder_path + "/WS_gfreq_uphase_N2000_NOMF_T20000_dt0.0100_nruns10_K%.3f_p=%.3f.tsv"%(Kvalues[i],j)
        #cols refers to timestep, avgmod, stdmod, avgphase,stdphase (of order parameter)
        df = pd.read_csv(filename,sep="\t",header=None)
        data.append([Kvalues[i],j,df])

# In[10]:

#data[0][x]#, x=0=> K, x=1=>p, x=2=> df

# In[11]:

def rinf_avg(df):
    lasts = df[1][int(T*.9):-1]
    return np.mean(lasts)

def rinf_std(df):
    lasts = df[1][int(T*.9):-1]
    return np.std(lasts)

# In[12]:

K_plot = []
p_plot = []
r_inf_avg = []
r_inf_std = []
for i in range(0,len(data)):
    t_plot = data[i][2][0]
    K_plot.append(data[i][0])
    p_plot.append(data[i][1])
    r_inf_avg.append(rinf_avg(data[i][2]))
    r_inf_std.append(rinf_std(data[i][2]))

# In[29]:

r_inf_mat = np.zeros(shape=[len(pvalues),len(Kvalues)])
for i in range(0,len(data)):
    r_inf_mat[find_nearest(pvalues, data[i][1])][find_nearest(Kvalues, data[i][0])] = rinf_avg(data[i][2])

plt.figure(figsize=[10,10])
ax = plt.gca()
#name = "Kuramoto oscillators on Watts-Strogatz network \n N = %d, r = %d, dt = %.3f, %s, n_runs = %d\n$r_{\\infty}$"%(N,2,dt,freq_plot, n_runs)

im = plt.imshow(r_inf_mat)
plt.title("Kuramoto oscillators on Watts-Strogatz network\n N=%d, $r_{WS}$=%d, T=%d, dt=%.3f, %s, n_runs=%d\n $r_{\\infty}$"%(N,3,T,dt,freq_plot,n_runs),fontsize=20)

plt.yticks(np.linspace(0,len(pvalues)-1,len(pvalues)),pvalues)
plt.ylabel("p",fontsize=18,rotation=0)
plt.xticks(np.linspace(0,len(Kvalues)-1,len(Kvalues)),Kvalues,rotation=45)
plt.xlabel("K",fontsize=18)

divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.15)
cbar = plt.colorbar(im, cax=cax)

plt.tight_layout()
plt.savefig(output_dir+config_name+"WS_rinf_heatmap.png")

plt.show()

# In[14]:

Kcs_df = pd.DataFrame([p_plot,K_plot,r_inf_avg,r_inf_std]).T

# In[15]:

Kcs_df = Kcs_df.sort_values(by=[0,1])

# In[31]:

Kc_plot = []

for i in range(0,len(pvalues)):

    idx_kc = find_nearest(Kcs_df[0+i*len(Kvalues)+1:len(Kvalues)*(i+1)][2],.5)
    rinf_value = find_nearest_value(Kcs_df[0+i*len(Kvalues)+1:len(Kvalues)*(i+1)][2],.5)
    df = Kcs_df[0+i*len(Kvalues)+1:len(Kvalues)*(i+1)]
    df = df.reset_index(inplace = False)
    Kc_plot.append([pvalues[i],df[1].iloc[idx_kc]])
    print("prob",pvalues[i],", rinf %.3f"%(rinf_value),", K",df[1].iloc[idx_kc])
Kc_plot = pd.DataFrame(Kc_plot, columns=["p","Kc"])

# In[54]:

import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# In[55]:

def func(x, a, b):
    return a+ b/x

# In[66]:

popt, pcov = curve_fit(func, Kc_plot["p"][2:],Kc_plot["Kc"][2:], p0=[1.6,2])
popt

# In[70]:

x_fit = np.linspace(Kc_plot["p"][2],1,100)
y_fit = func(x_fit,popt[0],popt[1])

# In[94]:

plt.figure(figsize=[14,6])
plt.grid(True,alpha=.3)
plt.title("Kuramoto oscillators on Watts-Strogatz network\n N=%d, $r_{WS}$=%d, T=%d, dt=%.3f, %s, n_runs=%d\n $K_{c}(p)$"%(N,3,T,dt,freq_plot,n_runs),fontsize=20)
plt.plot(x_fit,y_fit,label="y(p) = a + $\\frac{b}{p}$ fit")
plt.plot(Kc_plot["p"],Kc_plot["Kc"],c='b',marker='o',markersize=12,ls='',label="Raw Data")
plt.plot(Kc_plot["p"][2:],Kc_plot["Kc"][2:],c='r',marker='o',markersize=9,ls='--',linewidth=.7,label="Fitted Data")
plt.xticks(pvalues,pvalues)
plt.xlabel("p",fontsize=18,rotation=0)
plt.yticks(Kvalues,Kvalues)
plt.ylim(min(Kc_plot["Kc"])*.85,max(Kc_plot["Kc"])*1.1)
plt.legend(fontsize=18)

plt.text(.5,4,"a = %.4f$\pm$%.4f, b = %.4f$\pm$%.4f"%(popt[0],pcov[0,0],popt[1],pcov[1,1]),fontsize=20)

plt.tight_layout()
plt.savefig(output_dir+config_name+"WS_Kc(p).png")

import numpy,import scipy,from scipy,from sympy


# coding:utf-8
# usr/bin/python3
# python src/chapter32/chapter32note.py
# python3 src/chapter32/chapter32note.py
"""

Class Chapter32_1

Class Chapter32_2

Class Chapter32_3

Class Chapter32_4

"""
from __future__ import absolute_import, division, print_function

import math
import re
import numpy as np

if __name__ == '__main__':
    import stringmatch as sm
else:
    from . import stringmatch as sm

class Chapter32_1:
    """
    chapter32.1 note and function
    """
    def __init__(self):
        pass

    def note(self):
        """
        Summary
        ====
        Print chapter32.1 note

        Example
        ====
        ```python
        Chapter32_1().note()
        ```
        """
        print('chapter32.1 note as follow')
        print('第32章 字符串匹配')
        print('在文本编辑程序中,经常出现要在一段文本中找出某个模式的全部出现位置这一问题。典型情况是,一段文本是正在编辑的文件,',
            '所搜寻的模式是用户提供的一个特定单词。解决这个问题的有效算法能极大地提高文本编辑程序的响应性能',
            '字符串匹配算法也常常用于其他方面,例如在DNA序列中搜寻特定的模式')
        print('字符串匹配问题的形式定义是这样的:假设文本是一个长度为n的数组T[1..n],模式是一个长度为m<=n的数组P[1..m].',
            '进一步假设P和T的元素都是属于有限字母表∑表中的字符.例如可以有∑={0,1}或∑={a,b,...,z},字符数组P和T常称为字符串')
        print('如果0<=s<=n-m,并且T[S+1,...,s+m]=P[1..m](即对1<=j<=m,有T[s+j]=P[j]),则说模式P在文本T中出现且位移为s.',
            '(或者等价地,模式P在文本T中从位置s+1开始出现)。如果P在T中出现且位移为s,则称s为一个有效位移,否则称s为无效位移',
            '这样一来,字符串匹配问题就变成一个在一段指定的文本T中,找出某指定模式P出现所有有效位移的问题')
        print('本章的每个字符串匹配算法都对模式进行了一些预处理,然后找寻所有有效位移;我们称第二步为“匹配”.',
            '每个算法的总运行时间为预处理和匹配时间的总和.')
        print('32.2节介绍由Rabin和Karp发现的一种有趣的字符串匹配算法,该算法在最坏情况下的运行时间为Θ((n-m+1)m),虽然这一时间并不比朴素的算法好',
            '但是在平均情况和实际情况中,该算法的效果要好的多.这种算法也可以很好地推广到解决其他的模式匹配问题')
        print('32.3节中描述另一种字符串匹配算法,该算法构造一个特别设计的有限自动机,用来搜寻某给定模式P在文本中的出现的位置',
            '此算法用O(m|∑|)的预处理时间,但只用Θ(n)的匹配时间')
        print('32.4节介绍与其类似但更巧妙的Knuth-Morris-Pratt(或KMP)算法。该算法的匹配时间同样为Θ(n),但是将预处理时间降至Θ(m)')
        print('算法          预处理时间        匹配时间')
        print('朴素算法          0           O((n-m+1)m)')
        print('Rabin-Karp       Θ(m)        O((n-m+1)m)')
        print('有限自动机算法   O(m|∑|)         Θ(n)')
        print('KMP算法          Θ(m)           Θ(n)')
        print('记号与术语')
        print('  用∑*表示用字母表∑中的字符形成的所有有限长度的字符串的集合.在本章中仅考虑长度有限的字符串',
            '长度为0的空字符串用e表示,它也属于∑*.字符串x的长度用|x|表示.两个字符串x和y的连接表示为xy,其长度为|x|+|y|,由x的字符接y的字符组成')
        print('  如果对某个字符串与y∈∑*,有x=wy,就说字符串w是字符串x的前缀,表示为w∝x,得知|w|<=|x|.空字符串e既是每个字符串的前缀,也是每个字符串的后缀',
            '例如,有ab>abcca,cca<abcca.对任意字符串x和y及任意字符a,x>y当且仅当xa>ya.注意>和<都是传递关系')
        print('引理32.1(重叠后缀定理)假设x,y和z是满足x>z和y<z的三个字符串.如果|x|<=|y|,则x>y;如果|x|>=|y|,则y>x;如果|x|=|y|,则x=y')
        print('  本章中允许把比较两个等长的字符串是否相等的操作当做原语操作.如果对字符串的比较是从左往右进行,并且发现一个不匹配字符时比较就终止,',
            '则假设这样一个测试过程所需的时间是关于所发现的匹配字符数目的线性函数')
        print('32.1 朴素的字符串匹配算法')
        print('朴素的字符串匹配算法:它用一个循环来找出所有有效位移,该循环对n-m+1可能的每一个s值检查条件P[1..m]=T[s+1..s+m]')
        print('这种朴素的字符串匹配过程可以形象地看成用一个包含模式的“模板”沿文本滑动,同时对每个位移注意模板上的字符是否与文本的相应字符相等')
        print('NATIVE-STRING-MATCHER的运行时间为Θ((m-m+1)m)')
        print('在本章中还要介绍一种算法,它的最坏情况预处理时间为Θ(m),最坏情况匹配时间为Θ(n)')
        print('练习31.1-1 解答过程如下：')
        P = '0001'
        T = '000010001010001'
        sm.native_string_matcher(T, P)
        print('练习31.1-2 假设模式P中的所有字符都是不同的。试说明如何对一段n个字符的文本T加速过程NATIVE-STRING-MATCHER的执行速度,',
            '使其运行时间达到O(n)')
        print('练习31.1-3 假设模式P和文本T是长度分别为m和n的随机选取的字符串,其字符串属于d个元素的字母表∑={0,1,...,d-1},其中d>=2',
            '证明朴素算法第4行中隐含的循环所执行的字符比较的预计次数为')
        print('   (n-m+1)(1-d**-m)/(1-d**-1)<=2(n-m+1)')
        print('练习31.1-4 假设允许模式P中包含一个间隔字符◇,该字符可以与任意的字符串匹配(甚至可以与长度为0的字符串匹配)',
            '例如,模式ab◇ba◇c')
        # python src/chapter32/chapter32note.py
        # python3 src/chapter32/chapter32note.py

class Chapter32_2:
    """
    chapter32.2 note and function
    """
    def __init__(self):
        pass

    def note(self):
        """
        Summary
        ====
        Print chapter32.2 note

        Example
        ====
        ```python
        Chapter32_2().note()
        ```
        """
        print('chapter32.2 note as follow')
        print('Rabin-Karp算法')
        print('在实际应用中,Rabin和Karp所建议的字符串匹配算法能够较好的运行,我们还可以从中归纳出有关问题的其他算法,如二维模式匹配')
        print('Rabin-Karp算法预处理时间为')
        print('Θ(m),在最坏情况下的运行时间为O((n-m+1)m),但是它的平均情况运行时间还是比较好的')
        print('假定∑={0,1,2,...,9},这样每个字符都是一个十进制数字(一般情况下,可以假定每个字符都是基数为d的表示法中一个数字,d=|∑|).',
            '可以用一个长度为k的十进制数来表示由k个连续字符组成的字符串.因此,字符串31415就对应于十进制数31415',
            '如果输入字符既可以看做图形符号,也可以看做数字')
        print('已知一个模式P[1..m],设p表示其相应的十进制数的值。对于给定的文本T[1..n],用ts来表示其长度为m的子字符串T[s+1..s+m](s=0,1,...,n - m)相应的十进制数的值')
        print('当然,用ts来表示其长度为m的子字符串T[s+1..s+m](s=0,1,..,n-m)相应十进制数的值.当然,ts=p当且仅当T[s+1..s+m]=P[1..m],因此s是有效位移当且仅当ts=p',
            '如果能够在Θ(m)的时间内计算出p的值,并在总共Θ(n-m+1)的时间内计算出所有ts的值,那么通过把p值与每个ts值进行比较,能够在Θ(n)的时间内,求出有效位移s')
        print('可以运用霍纳法则,在Θ(m)的时间内计算出p的值：')
        print('  p=P[m]+10(P[m-1]+10(P[m-2]+...+10(P[2]+10P[1])...))')
        print('类似地,也可以在Θ(m)的时间内,根据T[1..m]计算出t0的值')
        print('为了在Θ(n-m)的时间内计算出剩余的值t1,t2,...,tn-m,可以在常数时间内根据ts计算出ts+1,这是因为霍纳法则')
        print('RABIN-KARP-MATCHER的预处理时间Θ(m),其匹配时间在最坏情况下为Θ((n-m+1)m),因为Rabin-Karp算法与朴素的字符串匹配算法一样,对每个有效位移进行显示验证',
            '如果P=a^m并且T=a^n,则验证所需的时间为Θ((n-m+1)m),因为n-m+1可能的位移中每一个都是有效位移')
        print('在许多实际作用中,有效位移数很少(如只有常数c个),因此,算法的期望匹配时间为O((n-m+1)+cm)=O(n+m),',
            '再加上处理伪命中点所需的时间.假设减少模q的值就像是从∑*到Zq上的一个随机映射,基于这种假设,进行启发性分析',
            '要正式证明这个假设是比较困难的,但是有一种可行的方法,就是假定q是从适当大的整数中随机得出的.可以预计伪命中的次数为O(n/q)',
            '因为可以估计出任意的ts对模q等价于p的概率为1/q.')
        print('Rabin-Karp算法的期望运行时间为:O(n)+O(m(v+n/q))')
        print('练习32.2-1 如果取模q=11,那么当Rabin-Karp匹配算法在文本T=3141592653589793中与搜寻模式P=26,会遇到多少个伪命中点')
        sm.native_string_matcher('3141592653589793', '26')
        sm.rabin_karp_matcher('3141592653589793', '26', 10, 11)
        print('练习32.2-2 如何扩展Rabin-Karp方法,使其能解决这样的问题:如何在文本字符串中搜寻出给定的k个模式中任何一个出现',
            '起初假定所有k个模式都是等长的.然后扩展算法允许不同长度的模式')
        print('练习32.2-3 试说明如何扩展Rabin-Karp方法以处理下列问题,在一个n*n二维字符串中搜寻出给定的m*m模式',
            '(可以使该模式在水平方向和垂直方向移动,但不可以把模式旋转)')
        print('练习32.2-4 Alice有一份很长的n位文件的复印件A=<an-1,an-2,...,a0>,Bob也有一份类似的文件B=<bn-1,bn-2,...,b0>',
            'Alice和Bob都希望知道他们的文件是否一样，为了避免传送整个文件A或B.',
            '运用下列快速概率检查手段,一起选择一个素数q>1000n,并从{0,1,...,q-1}中随机选取一个整数x,然后Alice求出：',
            'A(x)=(∑aixi) mod q)的值,Bob也用类似的方法计算出B(x).',
            '证明：如果A≠B,则A(x)=B(x)的概率至多为1/1000;如果两个文件相同,则A(x)的值必定等于B(x)的值')
        # python src/chapter32/chapter32note.py
        # python3 src/chapter32/chapter32note.py

class Chapter32_3:
    """
    chapter32.3 note and function
    """
    def __init__(self):
        pass

    def note(self):
        """
        Summary
        ====
        Print chapter32.3 note

        Example
        ====
        ```python
        Chapter32_3().note()
        ```
        """
        print('chapter32.3 note as follow')
        print('32.3 利用有限自动机进行字符串匹配')
        print('很多字符串匹配算法都要建立一个有限自动机,它通过对文本字符串T进行扫描的方法,',
            '找出模式P的所有出现位置.建立这样自动机的方法,用于字符串匹配：它们只对每个文本字符检查一次,并且检查每个文本字符的时间为常数',
            '因此,在建立好自动机后所需要的时间为Θ(n),但是如果∑很大,建立自动机所花的时间也可能是很多的')
        print('在本节的开头先定义有限自动机概念.考察一种一种特殊的字符串匹配自动机,并说明如何利用它找出一个模式在文本中的出现位置',
            '包括对一段给定的文本,如何模拟出字符串匹配自动机的执行步骤的一些细节.',
            '将说明对一个给定的输入模式,如何构造相应的字符串匹配自动机')
        print('有限自动机')
        print('  一个有限自动机M是一个5元组(Q,q0,A,∑,d)')
        print('   Q是一个状态的有限集合')
        print('   q0∈Q是初始状态')
        print('   A∈Q是一个接受状态集合')
        print('   ∑是有限的输入字母表')
        print('   d是一个从Q×∑到Q的函数,称为M的转移函数')
        print('  有限自动机开始于状态q0,每次读入输入字符串的一个字符.如果有限自动机在状态q时读入了输入字符a,则它从状态q变为状态d(q,a)(进行了一次转移).',
            '每当其状态q属于A时,就说自动机M接受了所有读入的字符串。没有被接收的输入称为被拒绝的输入')
        print('  有限自动机M可以推导出一个函数∮,称为终止函数,它是从∑*到Q的函数,并满足:∮(w)是M在扫描字符串w终止时的状态.',
            '因此,M接受字符串w当且仅当∮(w)∈A,函数∮由下列递归关系定义∮(e)∈q0',
            '∮(wa)=d(∮(w), a) 对于w∈∑*, a∈∑')
        print('字符串匹配自动机')
        print('  对每个模式P都存在一个字符串匹配自动机,必须在预处理阶段,根据模式构造出相应的自动机后,才能利用它来搜寻文本字符串',
            '关于模式P=ababaca的有限自动机的构造过程。从现在开始,假定P是一个已知的固定模式.为了使说明上的简洁,在下面的概念中将不特别指出对P的依赖关系')
        print('  为了详细说明与给定模式P[1..m]相应的字符串匹配自动机,首先定义一个辅助函数a,称为相应P的后缀函数。',
            '函数a是一个从∑*到{0,1,...,m}上定义的映射,a(x)是x的后缀P的最长前缀的长度：a(x)=max{k:Pk>x}')
        print('  为了清楚地说明字符串匹配自动机的操作过程,给出一个简单而有效的程序,用来模拟这样一个自动机(用它的变迁函数d来表示),在输入文本T[1..n]中,',
            '寻找长度为m的模式P的出现位置的过程,对于长度为m的模式的任意字符串匹配自动机来说,状态Q为{0,1,...,m},初始状态为0,唯一的接收态是状态m')
        print('  由FINITE-AUTOMATON-MATCHER的简单循环结构可以看出,对于一个长度为n的文本字符串,它的匹配时间为Θ(n)',
            '但是,这一匹配时间没有包括计算变迁函数d所需要的预处理时间.将在证明FINITE-AUTOMATON-MATCHER的正确性以后,再来讨论这一问题')
        print('  考察自动机在输入文本T[1..n]上进行的操作.将证明自动机扫过字符T[i]后,其状态为d(Ti).因为d(Ti)=m当且仅当P>Ti,',
            '所以自动机处于接收状态m,当且仅当模式P已经被扫描过,为了证明这个结论,要用到下面两条关于后缀函数o的引理')
        print('引理32.2 (后缀函数不等式) 对任意字符串x和字符a,有o(xa)>=o(x)+1')
        print('引理32.3 (后缀函数递归引理) 对任意x和字符a,如果q=o(x),则o(xa)=o(Pqa)')
        print('定理32.4 如果∮是字符串匹配自动机关于给定模式P的终态函数,T[1..n]是自动机的输入文本,对i=0,1,...,n,有∮(Ti)=o(Ti)')
        print('计算变迁函数')
        print('练习32.3-1 对模式P=aabab构造出相应的字符串匹配自动机,并说明它在文本字符串T=aaababaabaababaab上的操作过程')
        print(sm.finite_automaton_matcher('aaababaabaababaab', 10, 8))
        print('练习32.3-2 对字母表∑={a, b},画出与模式ababbabbababbababbabb相应的字符串匹配自动机状态转换图')
        print('练习32.3-3 如果由Pk>Pq蕴含着k=0或k=q,则称模式P是不可重叠的.试描述与不可重叠模式相应的字符串匹配自动机的状态转换图')
        print('练习32.3-4 已知两个模式P和P`,试描述如何构造一个有限自动机,使之能确定其中任意一个模式的所有出现位置.要求尽量使自动机的状态数最小')
        print('练习32.3-5 已知一个包括间隔字符的某模式P,说明如何构造一个有限自动机,使其在O(n)的时间内,找出P在文本T中的一次出现位置,其中n=|T|')
        # python src/chapter32/chapter32note.py
        # python3 src/chapter32/chapter32note.py

class Chapter32_4:
    """
    chapter32.4 note and function
    """
    def __init__(self):
        pass

    def note(self):
        """
        Summary
        ====
        Print chapter32.4 note

        Example
        ====
        ```python
        Chapter32_4().note()
        ```
        """
        print('chapter32.4 note as follow')
        print('32.4 Knuth-Morris-Pratt算法')
        print('Knuth、Morris和Pratt三人设计的线性时间字符串匹配算法。这个算法不用计算变迁函数d,匹配时间为Θ(n),只要用到辅助函数pi[1,m]',
            '它是在Θ(m)时间内,根据模式预先计算出来的.数组pi使得我们可以按需要,“现场”有效地计算(在平摊意义上来说)变迁函数d.',
            '粗略地说,对任意状态q=0,1,...,m和任意字符a∈∑,pi[q]的值包含了与a无关但在计算d(q,a)时需要的信息',
            '由于数组pi只有m个元素,而d有Θ(m|∑|个值,所以通过预先计算pi而不是d,使得时间减少了一个|∑|因子)')
        print('关于模式的前缀函数')
        print('  模式的前缀函数pi包含有模式与其自身的位移进行匹配的信息.这些信息可用于避免在朴素的字符串匹配算法中,对无位移进行测试,',
            '也可以避免在字符串匹配自动机中,对d的预先计算过程')
        print('KMP-MATCHER的大部分过程都是在模仿FINITE-AUTOMATON-MATCHER.KMP-MATCHER调用了一个辅助过程COMPUTE-PREFIX-FUNCTION来计算pi')
        print('运行时间分析')
        print('  运用平摊分析方法进行分析后可知,过程COMPUTE-PREFIX-FUNCTION的运行时间为Θ(m)')
        print('  在类似的平摊分析中,如果用q的值作为势函数,则KMP-MATCHER的匹配时间为Θ(n)',
            '与FINITE-AUTOMATON-MATCHER相比,通过运用pi而不是d,可使对模式进行预处理所需的时间由O(m|∑|)下降为Θ(m),同时保持实际的匹配时间为Θ(n)')
        print('前缀函数计算的正确性')
        print('  通过对前缀函数pi进行迭代,就能够列举出是某给定前缀Pq的后缀的所有前缀Pk,设',
            'pi*[q]={pi[q],pi(2)[q],pi(3)[q],...,pi(t)[q]}')
        print('引理32.5 (前缀函数迭代定理) 设P是长度为m的模式,其前缀函数为pi,对q=1,2,...,m,有pi*[q]={k:k<q且Pk>Pq}')
        print('引理32.6 设P是长度为m的模式,pi是P的前缀函数.对q=1,2,...,m,如果pi[q]>0,则pi[q]-1∈pi*[q-1]')
        print('推论32.7 设P是长度为m的模式,pi是P的前缀函数,对q=2,3,...,m')
        print('KMP算法的正确性')
        print('  过程KMP-MATCHER可以看做是过程FINITE-AUTOMATON-MATCHER的一次重新实现')
        print('练习32.4-1 当字母表为∑={a,b},计算相应于模式ababbabbabbababbabb的前缀函数pi')
        print('练习32.4-2 给出关于q的函数pi*[q]的规模的上界.举例说明所给出的上界是严格的')
        print('练习32.4-3 试说明如何通过检查字符串PT的pi函数,来确定模式P在文本T中的出现位置(由P和T并置形成的长度为m+n的字符串)')
        print('练习32.4-4 试说明如何通过以下方式对过程KMP-MATCHER进行改进:把第7行(不是第12行中)出现的pi替换为pi‘.对q=1,2,...,m的递归定义如下：')
        print('练习32.4-5 写出一个线性时间的算法,','以确定文本T是否是另一个字符串T‘的循环旋转,例如arc和car是彼此的循环旋转')
        print('练习32.4-6 给出一个有效的算法,计算出相应于某给定模式P的字符串匹配自动机的变迁函数d',
            '所给出的算法的运行时间应该是O(m|∑|).(提示：证明:如果q=m或P[q+1]!=a,则d(q,a)=d(pi[q],a))')
        print('思考题32-1 基于重复因子的字符串匹配')
        print('  设yi表示字符串y与其自身并置i次所得的结果.例如(ab)^3=ababab.如果对某个字符串y∈∑*和某个r>0有x=y^r,则称字符串x∈∑*具有重复因子r.',
            '设p(x)表示满足x具有重复因子r的最大值')
        print('  (a) 写出一个有效算法以计算出p(Pi)(i=1,2,...,m),算法的输入为模式P[1..m].算法的运行时间是多少？')
        print('  (b) 对任何模式p[1..m],设p(P)定义为max(1<=i<=m)p(Pi).证明：如果从长度为m的所有二进制字符串所组成的集中随机地选择模式P,则p*(P)的期望值是O(1)')
        print('  (c) 论证下列字符串匹配算法可以在O(p*(P)n+m)的运行时间内,正确地找出模式P在文本T[1..n]中的所有出现位置')
        # python src/chapter32/chapter32note.py
        # python3 src/chapter32/chapter32note.py

chapter32_1 = Chapter32_1()
chapter32_2 = Chapter32_2()
chapter32_3 = Chapter32_3()
chapter32_4 = Chapter32_4()

def printchapter32note():
    """
    print chapter32 note.
    """
    print('Run main : single chapter thirty-two!')
    chapter32_1.note()
    chapter32_2.note()
    chapter32_3.note()
    chapter32_4.note()

# python src/chapter32/chapter32note.py
# python3 src/chapter32/chapter32note.py

if __name__ == '__main__':
    printchapter32note()
else:
    pass


from IPython.display import HTML

HTML('''
The raw code for this IPython notebook is by default hidden for easier reading.
To toggle on/off the raw code, click here.''')

# From Ptolemy to Kepler

Ptolemy, Copernicus, Brahe, Kepler

# Cosmic Calculations 2.1: Kepler’s Third Law

First, let's review the three laws discovered by Kepler from the careful measurements of [Tycho Brache](https://physicsworld.com/a/kepler-and-tycho-brahe-the-odd-couple/).  This is one of the most ineteresting stories of scientific collaboration that transformed years of observations into laws about the universe.  I recommend to read/watch [***The Character of Physical Law***](https://www.youtube.com/watch?v=j3mhkYbznBk) by Richard Feynman if you want to indulge in the details.

By Both Előd at Hungarian Wikipedia, CC BY-SA 2.5, Link

## Kepler's laws:

***1. The orbit of every planet is an ellipse with the Sun at one of the two foci.***

In the figure below, you can imagine the yellow dot as the sun, and a planet would be moving on the blue curve a certain distance away from it.  The elliptical orbit can be described by the semi-major and semi-minor axes, which define the eccentricity of the orbit.

Look for the *perihelion* and *aphelion* of Earth's orbit. From those values, what's the flattening of Earth's orbit and its eccentricity?

import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets

# Set default font size for plots:
font = {'size'   : 18}
plt.rc('font',**font)

def elliptic_orbit(a,b,t):
    '''Plot an elliptical orbit and see the radial distance from one focal point
    a= semi-major axis
    b=semi-minor axis
    t=location at an angle between 0 and 360'''
    p=np.linspace(0,2*np.pi,360)
    x = a*np.cos(p)
    y = b*np.sin(p)
    plt.figure('Ellipse2',figsize=(10,5))
    plt.plot(x,y,'-')
    plt.axis('equal')
    plt.grid(True)
    #t is the angle varying from 0 to 360 degrees
    X = a*np.cos(t*np.pi/180)
    Y = b*np.sin(t*np.pi/180)
    #Conditionals in case of changing length of largest semi-major axis
    if a>=b:
        c=np.sqrt(a**2-b**2)
        plt.scatter(c,0,s=200,c='y')
        plt.scatter(-c, 0, s=200, facecolors='none', edgecolors='y')
        plt.arrow(c, 0, X-c, Y, head_width=0.1, head_length=0.1, fc='red', ec='red')
        plt.scatter(X,Y,s=50,c='b')
        f=(a-b)/a
        e=c/a
        #print('Orbital flattening : ',f)
        print('Orbital eccentricity : ',e)
        #plt.show()
    else:
        c=np.sqrt(b**2-a**2)
        plt.scatter(0,c,s=200,c='y')
        plt.scatter(0, -c, s=200, facecolors='none', edgecolors='y')
        plt.arrow(0, c, X, Y-c, head_width=0.1, head_length=0.1, fc='red', ec='red')
        plt.scatter(X,Y,s=50,c='b')
        f=(b-a)/b
        e=c/b
        #print('Orbital flattening : ',f)
        print('Orbital eccentricity : ',e)
    plt.show()
    return

interactive(elliptic_orbit, a = (0,20,1),b=(0,20,1),t=(0,360,20),continuous_update=False)

The reason the orbits are ellipses, and not circles, would not be understood until the arrival of Newton's equations on Gravitation.

***2. A line joining a planet and the Sun sweeps out equal areas during equal intervals of time.***

By Gonfer (talk) - Gonfer, CC BY-SA 3.0, Link

When Kepler discovered his third law ($p^2 = a^3$), he knew only that it applied to the orbits of planets about the Sun. In fact, it applies to any orbiting object as long as the following two conditions are met:

1. The object orbits the Sun or another star of precisely the same mass.
2. We use units of years for the orbital period and AU for the orbital distance. (Newton extended the law to all orbiting objects; see [Cosmic Calculations 7.1](#CC-7.1).)

In other words, these two conditions make the relationship a perfect equality.

**Example 1:** The largest asteroid, Ceres, orbits the Sun at an average distance (semimajor axis) of 2.77 AU. What is its orbital period?

***Solution:*** Both conditions are met, so we solve Kepler’s third law for the orbital period $p$ and substitute the given orbital distance, $a = 2.77~AU$.

$$p^2 = a^3$$

$$ p = \sqrt{a^3} = \sqrt{2.77^3} \approx 4.6~y$$

Ceres has an orbital period of 4.6 years.

**Example 2:** A planet is discovered orbiting every three months around a star of the same mass as our Sun. What is the planet’s average orbital distance?

***Solution:*** The ﬁrst condition is met, and we can satisfy the second by converting the orbital period from months to years: $p = 3$ months = 0.25 year. We now solve Kepler’s third law for the average distance a:

$a = \sqrt[3]{p^2}$

$a = \sqrt[3]{0.25^2} \approx 0.40~AU$

The planet orbits its star at an average distance of $0.40~AU$, which is nearly the same as Mercury’s average distance from the Sun.

These observations offered clear proof that Earth is not the center of everything.* Although we now recognize that Galileo won the day, the story was more complex in his own time, when Catholic Church doctrine still held Earth to be the center of the universe. On June 22, 1633, Galileo was brought before a Church inquisition in Rome and ordered to recant his claim that Earth orbits the Sun. Nearly 70 years old and fearing for his life, Galileo did as ordered and his life was spared. However, legend has it that as he rose from his knees, he whispered under his breath, Eppur si muove— Italian for “And yet it moves.” (Given the likely consequences if Church ofﬁcials had heard him say this, most historians doubt the legend.) The Church did not formally vindicate Galileo until 1992, but the Church had given up the argument long before that. Today, Catholic scientists are at the forefront of much astronomical research, and ofﬁcial Church teachings are compatible not only with Earth’s planetary status but also with the theories of the Big Bang and the subsequent evolution of the cosmos and of life.

<object classid="clsid:D27CDB6E-AE6D-11cf-96B8-444553540000" width="750" height="400"><param name="movie" value="KeplerFirstLaw.swf" /><object type="application/x-shockwave-flash" data="KeplerFirstLaw.swf" width="750" height="400">
flash animation
</object></object>

[Go back to the top of the page](#top-title)