import argparse
import csv
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import re
import numpy as bp
import scipy
import pysam
import multiprocessing
import gzip

class RecessiveModel:
	def __init__(self, af = 1e-2, DeepVariant = True, SBPV_cutoff=1e-5, DP_cutoff=7, AB_cutoff1=0.2, AB_cutoff2=0.8):
		self.SBPV_cutoff = SBPV_cutoff
		self.DP_cutoff_T1 = DP_cutoff
		self.AB_cutoff1_T1 = AB_cutoff1
		self.AB_cutoff2_T1 = AB_cutoff2
		self.DP_cutoff_T2 = 20
		self.AB_cutoff_T2 = 0.3
		self.LGD = set(["splice_acceptor_variant", "splice_donor_variant", "stop_gained",
			"stop_lost", "start_lost", "frameshift_variant"])
		self.CADD_cutoff = 25
		self.REVEL_cutoff = 0.5
		self.AF_cutoff = af
		#self.C = ["syn","lgd","dmis","lgd/dmis"]
		self.C = ["syn","lgd","mis","cadd15","cadd20","cadd25","revel.5","mvp2.85","mpc1","lgd_cadd25"]
		self.DeepVariantFilter = DeepVariant
		return

	# compute NHet, NHom, AC and AF of each site, within a population, Write to a VCF file.
	def ComputeSiteAF(self, InpVCF, Indvs, prefix, OutVCF):
		fin = gzip.open(InpVCF, 'rt')
		fout = open(OutVCF, 'w')
		for l in fin:
			if l.startswith("##"):
				fout.write(l)
				continue
			elif l.startswith("#"):
				llist = l.strip().split("\t")
				#Head = llist[]
				indvs = llist[9:]
				index = self.GetIndex(Indvs, indvs)
				fout.write("##ComputeAF={}\n".format(prefix))
				fout.write("\t".join(llist[:9])+"\n")
			else:
				llist = l.strip().split("\t")
				length = (len(llist[4].split(","))+1)
				AFs = [0] * length
				ACs = [0] * length
				AC_Het = [0] * length
				AC_Hom = [0] * length
				AN = 0
				GTs = llist[9:]
				for idx in index:
					GT = self.GenotypeQC(llist[8], GTs[idx])
					if GT:
						A1, A2 = GT[0], GT[1]
						ACs[A1] += 1
						ACs[A2] += 1
						if A1 != A2:
							AC_Het[A2] += 1
						else:
							AC_Hom[A2] += 1
						AN += 2
				for i in range(length):
					try:
						AFs[i] = str(float(ACs[i])/AN)
					except ZeroDivisionError:
						AFs[i] = "0"
					ACs[i] = str(ACs[i])
					AC_Het[i] = str(AC_Het[i])
					AC_Hom[i] = str(AC_Hom[i])
				New = ";{}_AN={};{}_AC={};{}_AF={};{}_AC_Het={};{}_AC_Hom={}".format(prefix, AN, prefix, ",".join(ACs[1:]), prefix, ",".join(AFs[1:]), prefix, ",".join(AC_Het[1:]), prefix, ",".join(AC_Hom[1:]))
				llist[7] = llist[7] + New
				fout.write("\t".join(llist[:8])+"\n")
		return
	def GetIndex(self, Indvs, indvs):
		res = []
		for indv in Indvs:
			try:
				res.append(indvs.index(indv))
			except:
				continue
		return res

	def isRef(self, GT):
		if GT[0] == GT[1] and GT[0] == 0:
			return True
		else:
			return False

	def isHet(self, GT):
		if GT[0] == 0:
			return True
		else:
			return False

	def isHom(self, GT):
		if GT[0] != 0 and GT[1] != 0:
			return True
		else:
			return False

	# Return genotype as [A1, A2] if pass QC , return False otherwise
	def GenotypeQC(self, SPID, fmt, gt_dat):
		# Part I: result of GATK filter; Part II: result of DeepVariant Filter
		FLAG_TIER0 = 0; FLAG_TIER1 = 0; FLAG_TIER2 = 0
		tmp = {}
		for k,v in zip(fmt.split(":"), gt_dat.split(":")):
			tmp[k] = v
		GT = tmp["GT"].split("/")
		if GT[0] != "." and GT[1] != ".":
			GT = [int(i) for i in GT]
		else: # Missing GT
			print("Missing GT (GT)", gt_dat)
			return None, None
		if tmp["GQ"] == ".":
			#return False
			print("Missing GT (GQ)", gt_dat)
			return [0,0], None
		#elif float(tmp["GQ"]) < 60:
		#	return False
		#if GT[0] == "." or GT[1] == ".":
		#	return False
		if tmp["GT"] != "0/0":
			if self.isHet(GT):
				if float(tmp["DP"]) < self.DP_cutoff:
					return False
				if float(tmp["AD"].split(",")[1])/float(tmp["DP"]) < self.AB_cutoff1: # or float(tmp["AD"].split(",")[1])/float(tmp["DP"]) > self.AB_cutoff2:
					return False
			else:
				if not self.isHom(GT):
					print('Error detect Hom', GT, gt_dat)
		else:
			return GT, True

	def GenotypeQC_old(self, fmt, gt_dat):
		tmp = {}
		for k,v in zip(fmt.split(":"), gt_dat.split(":")):
			tmp[k] = v
		GT = tmp["GT"].split("/")
		if tmp["GQ"] == ".":
			#return False
			return [0,0]
		elif float(tmp["GQ"]) < 60:
			return False
		if GT[0] == "." or GT[1] == ".":
			return False
		if tmp["GT"] != "0/0":
			#print(fmt, gt_dat)
			#if "." in tmp["SBPV"]:
			#	pass
			#elif float(tmp["SBPV"]) < self.SBPV_cutoff:
			#	return False
			if float(tmp["DP"]) < self.DP_cutoff:
				return False
			if float(tmp["AD"].split(",")[1])/float(tmp["DP"]) < self.AB_cutoff1: # or float(tmp["AD"].split(",")[1])/float(tmp["DP"]) > self.AB_cutoff2:
				return False
		return [int(i) for i in GT]

	def LoadPedigree2(self, PedFil, Samples):
		PedFil = "/home/local/users/jw/Genetics_Projects/SPARK/spark_genomics/dat/EUR_Trios.ped"
		Trios = []
		reader = csv.reader(open(PedFil, 'rt'), delimiter="\t")
		counter = 0
		for row in reader:
			row.append(Samples.index(row[1]))
			counter += 1
			if counter == 1:
				tmp = Family(row[0])
				tmp.Proband = Sample(row)
			if counter == 2:
				tmp.Father = Sample(row)
			if counter == 3:
				tmp.Mother = Sample(row)
				Trios.append(tmp)
				counter = 0
		return Trios
	def LoadPedigree(self, PedFil, Samples):
		#PedFil = "/home/local/users/jw/Genetics_Projects/SPARK/spark_genomics/dat/EUR_Fams.ped"
		PedFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/recessive/EUR_Fams.ped"
		Fams = []
		Indvs = []
		reader = csv.reader(open(PedFil, 'rt'), delimiter="\t")
		PreFamID, tmp = None, None
		for row in reader:
			row.append(Samples.index(row[1])) # Add sample index in VCF header, to locate genotype
			FamID = row[0]
			Indvs.append(row[1])
			if FamID != PreFamID:
				if tmp != None:
					Fams.append(tmp)
				tmp = Family(FamID)
				PreFamID = FamID
				tmp.Proband = Sample(row)
			else:
				if row[1] == tmp.Proband.Father:
					tmp.Father = Sample(row)
				elif row[1] == tmp.Proband.Mother:
					tmp.Mother = Sample(row)
				else:
					tmp.Siblings.append(Sample(row))
		Indvs = set(Indvs)
		return Fams, Indvs

	def getINFO(self, info_string):
		infolist = info_string.split(';')
		infodict = {}
		for kv in infolist:
			kv = kv.split('=')
			if len(kv) == 2:
				k, v = kv
				infodict[k] = v
		return infodict

	def match_allele_csq(self, Ref, Alts, csq_head, csq_string):
		# Trim Leading Base
		Alts = Alts.split(",")
		if len(list(set([x[0] for x in Alts])))==1 and Ref[0] == list(set([x[0] for x in Alts]))[0]:
			_Ref = Ref[1:] if len(Ref[1:]) >0 else "-"
			_Alts = [Alt[1:] if len(Alt[1:]) >0 else "-" for Alt in Alts]
			#print (Ref, Alts, ";", _Ref, _Alts)
		else:
			_Alts = Alts
		res = {}
		csqs = csq_string.split(",")
		csqs = [dict(zip(csq_head, vep.split("|"))) for vep in csqs]
		for i, Alt in enumerate(Alts):
			res[Alt] = []
			for j, csq in enumerate(csqs):
				if csq["Allele"] == _Alts[i]:
					csq["Consequence"] = csq["Consequence"] .split("&")
					res[Alt].append(csq)
		return res

	def search_severe_consequence(self, var_k, Allele_CSQ_dict, Alt):
		severe_consequence = None
		severe_trans = None
		severe_idx = 0
		for i in range(len(Allele_CSQ_dict[Alt])):
			consequence = Allele_CSQ_dict[Alt][i]["Consequence"]
			Transcript = Allele_CSQ_dict[Alt][i]["Feature"]
			#print(Allele_CSQ_dict[Alt][i]["BIOTYPE"])
			if Allele_CSQ_dict[Alt][i]["BIOTYPE"] != "protein_coding":
				continue
			if len(set(consequence).intersection(self.LGD))>= 1:
				return i, consequence, Transcript
			elif consequence[0] == "missense_variant":
				return i, consequence, Transcript
			elif consequence[0] == "synonymous_variant":
				severe_consequence = consequence
				severe_trans = Transcript
				severe_idx = i
		if severe_consequence == None:
			return 0, "non-coding", None
		else:
			return 0, severe_consequence, severe_trans

	#def Recessive(self, Chr, GenotypeFil, VEPFil, AFFil, GenecodeFil):
	def Recessive(self, Chr, GenotypeFil, VEPFil, GenecodeFil):
		GenotypeFil = pysam.TabixFile(GenotypeFil)
		VEPFil = pysam.TabixFile(VEPFil)
		#AFFil = pysam.TabixFile(AFFil)
		Genes, Transtripts = LoadGeneCode(GenecodeFil)
		#CSQ_header = [X.strip().split("Format: ")[1].rstrip('>\"').split("|") for X in VEPFil.header if X.startswith("##INFO=<ID=CSQ")][0]
		CSQ_header = [X.strip().split("Format: ")[1].rstrip('>\"').split("|") for X in GenotypeFil.header if X.startswith("##INFO=<ID=CSQ")][0]
		#Samples = GenotypeFil.header[-1].split("\t")[9:]
		Samples = VEPFil.header[-1].split("\t")[9:]
		OutFil = open("Rec.FamTest.Chr{}.txt".format(Chr), 'w')
		OutFil.write("#Gene\t" + "\t".join(["{}.obs\t{}.haps".format(t,t) for t in self.C]) + "\n")
		OutFil2 = open("Rec.FamTest.Chr{}.sup.txt".format(Chr), 'w')
		OutFil2.write("#Gene\t" + "\t".join(["{}.NCantPhase\t{}.CantPhaseFams\t{}.N3vars".format(t,t,t) for t in self.C]) + "\n")
		#OutFil2.write("{}\t{}\t{}\n".format(Gene, "\t".join("{}\t{}\t{}".format(CantPhase[t], CantPhase_fams[t], MoreThanThree[t]) for t in self.C)))
		Trios = self.LoadPedigree("a", Samples)
		if self.DeepVariantFilter:
			self.DeepVarTBX = LoadDeepVar(Samples)
		for i, (Gene,GTF) in enumerate(Genes.items()): # iterate through genes
			Gene_Fam_dat = {} # store genotypes for each fam, group by variant categories
			for cat in self.C:
				Gene_Fam_dat[cat] = {}
				for i, trio in enumerate(Trios):
					Gene_Fam_dat[cat][trio.FamID] = []
			start, end = int(GTF.start), int(GTF.end)
			veps, cohort, genotypes = [],[], []
			for term in VEPFil.fetch(Chr, start, end):
				veps.append(term)
			#for term in AFFil.fetch(Chr, start, end):
			#	cohort.append(term)
			for term in GenotypeFil.fetch(Chr, start, end):
				genotypes.append(term)
			#for var in zip(veps, cohort, genotypes):
			for var in zip(veps, genotypes):
				llist = var[0].split("\t")
				llist2 = var[1].split("\t")
				Chr, Pos, Ref, Alts = llist[0], llist[1], llist[3], llist[4]
				#cohort_af = list(map(float, self.getINFO(var[1].split("\t")[7])["EUR_AF"].split(",")))
				fmt = llist2[8]
				Sample_genotypes = llist2[9:]
				infodict = self.getINFO(llist[7])
				mappability = float(infodict.get("Mappability", 1))
				if mappability != 1:
					continue
				Allele_CSQ_dict = self.match_allele_csq(Ref, Alts, CSQ_header, infodict["CSQ"])
				for i, Alt in enumerate(Alts.split(",")):
					var_k = "{}:{}:{}:{}".format(Chr, Pos, Ref, Alt)
					try:
						Allele_CSQ_dict[Alt][0]["gnomADg_AF_NFE"] = Allele_CSQ_dict[Alt][0]["gnomADg_AF_NFE"].split("&")[0]
						Allele_CSQ_dict[Alt][0]["gnomADe_AF_NFE"] = Allele_CSQ_dict[Alt][0]["gnomADe_AF_NFE"].split("&")[0]
						vep = Allele_CSQ_dict[Alt][0]
						gnomADg_af = 0 if (vep["gnomADg_AF_NFE"] == "" or vep["gnomADg_AF_NFE"] == ".")\
								else float(vep["gnomADg_AF_NFE"])
						gnomADe_af = 0 if (vep["gnomADe_AF_NFE"] == "" or vep["gnomADe_AF_NFE"] == ".")\
								else float(vep["gnomADe_AF_NFE"])
						#af = cohort_af[i]
						#if gnomADg_af > 1e-2 or af > 1e-2 or af == 0:
						#if max(gnomADg_af, af) > 1e-2 or af == 0:
						#if max(gnomADg_af, af) > self.AF_cutoff or af == 0:
						if gnomADg_af > self.AF_cutoff:
							continue
						# cons = Allele_CSQ_dict[Alt][0]["Consequence"]
						idx_anno, cons, trans = self.search_severe_consequence(var_k, Allele_CSQ_dict, Alt)
						#print (cons)
						if len(set(cons).intersection(self.LGD))>= 1:
							#print (gnomADe_af, af, cons)
							#self.LookUpBiallic(i, "lgd", fmt, Sample_genotypes, Trios)
							Gene_Fam_dat = self.AddVar(i, var_k, "lgd", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "synonymous_variant" in set(cons):
							Gene_Fam_dat = self.AddVar(i, var_k, "syn", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons):
							Gene_Fam_dat = self.AddVar(i, var_k, "mis", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons) and float(Allele_CSQ_dict[Alt][0]["CADD_PHRED"]) > 15:
							Gene_Fam_dat = self.AddVar(i, var_k, "cadd15", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons) and float(Allele_CSQ_dict[Alt][0]["CADD_PHRED"]) > 20:
							Gene_Fam_dat = self.AddVar(i, var_k, "cadd20", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons) and float(Allele_CSQ_dict[Alt][0]["CADD_PHRED"]) > 25:
							Gene_Fam_dat = self.AddVar(i, var_k, "cadd25", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if ("missense_variant" in set(cons) and float(Allele_CSQ_dict[Alt][0]["CADD_PHRED"]) > 25) or (len(set(cons).intersection(self.LGD))>= 1) :
							Gene_Fam_dat = self.AddVar(i, var_k, "lgd_cadd25", fmt, Sample_genotypes, Trios, Gene_Fam_dat)

						revel = 0 if Allele_CSQ_dict[Alt][0]["REVEL_score"].split("&")[0] in ["", "."] else float(Allele_CSQ_dict[Alt][0]["REVEL_score"].split("&")[0])
						mpc = 0 if Allele_CSQ_dict[Alt][0]["MPC_score"].split("&")[0]  in ["", "."]  else float(Allele_CSQ_dict[Alt][0]["MPC_score"].split("&")[0])
						mvp2 = 0 if Allele_CSQ_dict[Alt][0]["MVP2_rankscore"].split("&")[0] in ["", "."]  else float(Allele_CSQ_dict[Alt][0]["MVP2_rankscore"].split("&")[0])
						if "missense_variant" in set(cons) and revel > 0.5:
							Gene_Fam_dat = self.AddVar(i, var_k, "revel.5", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons) and mpc > 1:
							Gene_Fam_dat = self.AddVar(i, var_k, "mpc1", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
						if "missense_variant" in set(cons) and mvp2 > 0.85:
							Gene_Fam_dat = self.AddVar(i, var_k, "mvp2.85", fmt, Sample_genotypes, Trios, Gene_Fam_dat)
					except KeyError as e:
						print(e)
						print("KeyError", Ref, Alts, Alt, Allele_CSQ_dict)
						return
					except IndexError:
						print("IndexError", Ref, Alts, llist[7], Allele_CSQ_dict)
						return
			res = self.Phasing_N_Count(Gene_Fam_dat, Trios)
			OBS = {}
			EXP = {}
			CantPhase = {}
			MoreThanThree = {}
			CantPhase_fams = {}
			for t in self.C:
				OBS[t] = (res[t][0] + res[t][1])
				EXP[t] = (res[t][2])
				CantPhase[t] = res[t][3]
				CantPhase_fams[t] = ",".join(res[t][4])
				MoreThanThree[t] = res[t][5]
			OutFil.write("{}\t{}\n".format(Gene, "\t".join("{}\t{}".format(OBS[t], EXP[t]) for t in self.C)))
			OutFil2.write("{}\t{}\n".format(Gene, "\t".join("{}\t{}\t{}".format(CantPhase[t], CantPhase_fams[t], MoreThanThree[t]) for t in self.C)))
		return

	def AddVar(self, i, var_k, Vartype, fmt, gts, Trios, Gene_Fam_dat):
		N_mendelian_Error = 0
		trio_var_pairs = []
		for j, trio in enumerate(Trios):
			if N_mendelian_Error >= 2: # drop the site if >2 fam with mendelian error
				return Gene_Fam_dat
			prob, fa, mo, sibs = trio.Proband, trio.Father, trio.Mother, trio.Siblings

			gt_prob, gt_fa, gt_mo = self.GenotypeQC(fmt, gts[prob.index]), self.GenotypeQC(fmt, gts[fa.index]), self.GenotypeQC(fmt, gts[mo.index])
			gt_sibs = [self.GenotypeQC(fmt, gts[x.index]) for x in sibs]

			if gt_prob == False or gt_fa == False or gt_mo == False:
				continue # Failed QC
			elif ( (gt_prob[1] not in [0, i+1]) or (gt_fa[1] not in [0, i+1]) or (gt_mo[1] not in [0, i+1]) ) or (gt_prob[1] == 0 and gt_fa[1] == 0 and gt_mo[1] == 0):
				continue # Not this allele
			sib_fail_qc = False
			for gt in gt_sibs:
				if gt == False:
					sib_fail_qc = True
			if sib_fail_qc:
				continue
			elif (gt_prob[0] not in gt_fa or gt_prob[1] not in gt_mo) and (gt_prob[1] not in gt_fa or gt_prob[0] not in gt_mo):
				N_mendelian_Error += 1
				continue # Mendelian Error
			else:
				gt_prob, gt_fa, gt_mo = self.gt_recode(gt_prob), self.gt_recode(gt_fa), self.gt_recode(gt_mo)
				gt_sibs = [self.gt_recode(gt) for gt in gt_sibs]
				#Gene_Fam_dat[trio.FamID][Vartype].append([var_k, gt_prob, gt_fa, gt_mo])
				#Gene_Fam_dat[Vartype][trio.FamID].append([var_k, gt_prob, gt_fa, gt_mo, gt_sibs])
				trio_var_pairs.append((trio.FamID, [var_k, gt_prob, gt_fa, gt_mo, gt_sibs]))
		for FamID, dat in trio_var_pairs:
			Gene_Fam_dat[Vartype][FamID].append(dat)
		return Gene_Fam_dat

	def gt_recode(self, gt):
		if gt[0] != 0 :
			gt[0] = 1
		if gt[1] != 0 :
			gt[1] = 1
		return gt

	def Phasing_N_Count(self, Gene_Fam_dat, Trios):
		res = {}
		for t in self.C:
			N_hom = 0
			N_chet = 0
			N_hom_chet = 0
			N_haps = 0
			N_cant_phase = 0
			cant_phase_fam = []
			N_more_than_three = 0
			for i, trio in enumerate(Trios):
				variants_in_fam = Gene_Fam_dat[t][trio.FamID] #list of variants in this gene in this fam
				#for item in variants_in_fam
				if len(variants_in_fam) == 1: #only 1 variant
					var_k, gt_pro, gt_fa, gt_mo, gt_sibs = variants_in_fam[0]
					N_haps += sum(gt_fa + gt_mo)
					for gt in [gt_pro] + gt_sibs:
						if gt == [1,1]:
							N_hom += 1
				elif len(variants_in_fam) == 2: # 2 variants
					v1, gt_p1, gt_f1, gt_m1, gt_sibs1 = variants_in_fam[0]
					v2, gt_p2, gt_f2, gt_m2, gt_sibs2 = variants_in_fam[1]
					gts1 = [gt_p1] + gt_sibs1
					gts2 = [gt_p2] + gt_sibs2
					if (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,0]) or (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,0]):
						# 0/0 0/1 -> 0/1
						# 0/1 0/0 -> 0/1
						N_haps += 2
						for gt1, gt2 in zip(gts1 ,gts2):
							if gt1 == [0,1] and gt2 == [0,1]:
								N_chet += 1
					elif (gt_f1 == [0,1] and gt_m1 == [0,1] and gt_p1 == [0,1]) or (gt_f2 == [0,1] and gt_m2 == [0,1] and gt_p2 == [0,1]):
						# Unable to phase
						N_cant_phase += 1
						N_haps += 4
						#N_chet += 1
						cant_phase_fam.append(trio.FamID)
						for gt1, gt2 in zip(gts1, gts2):
							if gt1 == [0,1] and gt2 == [0,1]:
								N_chet += 1
				elif len(variants_in_fam) >= 2: # more than 2 variants
					N_more_than_three += 1
			res[t] = (N_hom, N_chet, N_haps, N_cant_phase, cant_phase_fam, N_more_than_three)
		return res

	def Phasing_N_Count2(self, Gene_Fam_dat, Trios):
		res = {}
		for t in self.C:
			N_hom = 0
			N_chet = 0
			N_hom_chet = 0
			N_haps = 0
			N_cant_phase = 0
			cant_phase_fam = []
			N_more_than_three = 0
			for i, trio in enumerate(Trios):
				variants_in_fam = Gene_Fam_dat[t][trio.FamID] #list of variants in this gene in this fam
				#for item in variants_in_fam
				if len(variants_in_fam) == 1: #only 1 variant
					var_k, gt_pro, gt_fa, gt_mo = variants_in_fam[0]
					N_haps += sum(gt_fa + gt_mo)
					if gt_pro == [1,1]:
						N_hom += 1
				elif len(variants_in_fam) == 2: # 2 variants
					v1, gt_p1, gt_f1, gt_m1 = variants_in_fam[0]
					v2, gt_p2, gt_f2, gt_m2 = variants_in_fam[1]
					if (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,0]) or (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,0]):
						# 0/0 0/1 -> 0/1
						# 0/1 0/0 -> 0/1
						N_haps += 2
						if gt_p1 == [0,1] and gt_p2 == [0,1]:
							N_chet += 1
						#elif (gt_p1 == [0,0] and gt_p2 == [0,1]) or (gt_p1 == [0,1] and gt_p2 == [0,0]) or (gt_p1 == [0,0] and gt_p2 == [0,0]):
						#	N_chet += 0
					#elif (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,1]):
					#	if gt_p1 == [0,1] and gt_p2 == [0,0]:
					#		Nhaps += 3
					#elif (gt_f1 == [0,0] and gt_m1 == [0,1] and gt_f2 == [0,1] and gt_m2 == [0,1]):
					#	if gt_p1 == [0,1] and gt_p2 == [0,0]:
					#		Nhaps += 3
					#elif (gt_f1 == [0,1] and gt_m1 == [0,1] and gt_f2 == [0,0] and gt_m2 == [0,1]):
					#	if gt_p1 == [0,0] and gt_p2 == [0,1]:
					#		Nhaps += 3
					#elif (gt_f1 == [0,1] and gt_m1 == [0,0] and gt_f2 == [0,1] and gt_m2 == [0,0]):
					#	if gt_p1 == [0,0] and gt_p2 == [0,1]:
					#		Nhaps += 3
					elif (gt_f1 == [0,1] and gt_m1 == [0,1] and gt_p1 == [0,1]) or (gt_f2 == [0,1] and gt_m2 == [0,1] and gt_p2 == [0,1]):
						# Unable to phase
						N_cant_phase += 1
						cant_phase_fam.append(trio.FamID)
				elif len(variants_in_fam) >= 2: # more than 2 variants
					N_more_than_three += 1
			res[t] = (N_hom, N_chet, N_haps, N_cant_phase, cant_phase_fam, N_more_than_three)
		return res

	def LookUpBiallic(self, i, Vartype, fmt, gts, Trios):
		for j, trio in enumerate(Trios):
			prob, fa, mo = trio.Proband, trio.Father, trio.Mother
			#print(fmt, gts[prob.index])
			gt_prob, gt_fa, gt_mo = self.GenotypeQC(fmt, gts[prob.index]), self.GenotypeQC(fmt, gts[fa.index]), self.GenotypeQC(fmt, gts[mo.index])
			if gt_prob == False or gt_fa == False or gt_mo == False:
				continue # gt failed QC
			elif (gt_prob[0] not in gt_fa and gt_prob[1] not in gt_mo) or (gt_prob[1] not in gt_fa and gt_prob[0] not in gt_mo):
				continue # mendelian error
			else:
				# Phasing
				if gt_prob[1] == i+1 and gt_prob[0] == i+1: # Hom
					Trios[j].pro_haps[Vartype] = [1,1]
					#print("12", Trios[j].pro_haps[Vartype], gt_prob, gt_fa, gt_mo)
					if gt_fa[0] == i+1 and gt_fa[1] == i+1:
						Trios[j].fa_haps[Vartype] = [1,1]
					else:
						Trios[j].fa_haps[Vartype][0] = 1
					if gt_mo[0] == i+1 and gt_mo[1] == i+1:
						Trios[j].mo_haps[Vartype] = [1,1]
					else:
						Trios[j].mo_haps[Vartype][0] = 1

				elif gt_prob[1] == gt_fa[1] and gt_mo[1] == 0 and gt_prob[1] == i+1 : #paternal transmitted
					if gt_fa[0] == i+1: # transmitted from hom parternal
						Trios[j].fa_haps[Vartype] = [1,1]
					else: # transmitted from het paternal
						Trios[j].fa_haps[Vartype][0] = 1
					Trios[j].pro_haps[Vartype][0] = 1

				elif gt_prob[1] == gt_mo[1] and gt_fa[1] == 0 and gt_prob[1] == i+1: #maternal transmitted
					if gt_mo[0] == i+1: # transmitted from hom maternal
						Trios[j].mo_haps[Vartype] = [1,1]
					else: # transmitted from het maternal
						Trios[j].mo_haps[Vartype][0] = i+1
					Trios[j].pro_haps[Vartype][1] = 1

				elif gt_prob[1] == 0: # proband 0/0
					if gt_fa[1] == 1: # father has one hap
						if Trios[j].fa_haps[Vartype][0] == 1:
							Trios[j].fa_haps[Vartype][1] = 0
						elif Trios[j].fa_haps[Vartype][0] == 0:
							Trios[j].fa_haps[Vartype][1] = 1
					if gt_mo[1] == 1: # mother has one hap
						if Trios[j].mo_haps[Vartype][0] == 1:
							Trios[j].mo_haps[Vartype][1] = 0
						elif Trios[j].mo_haps[Vartype][0] == 0:
							Trios[j].mo_haps[Vartype][1] = 1

	def LookUpBiallicLGD_DMIS(self, i, Vartype, fmt, gts, Trios):
		for j, trio in enumerate(Trios):
			prob, fa, mo = trio.Proband, trio.Father, trio.Mother
			#print(fmt, gts[prob.index])
			gt_prob, gt_fa, gt_mo = self.GenotypeQC(fmt, gts[prob.index]), self.GenotypeQC(fmt, gts[fa.index]), self.GenotypeQC(fmt, gts[mo.index])
			if gt_prob == False or gt_fa == False or gt_mo == False:
				continue # gt failed QC
			elif (gt_prob[0] not in gt_fa and gt_prob[1] not in gt_mo) or (gt_prob[1] not in gt_fa and gt_prob[0] not in gt_mo):
				continue # mendelian error
			else:
				# Phasing
				if gt_prob[1] == i+1 and gt_prob[0] == i+1: # Hom
					Trios[j].pro_haps[Vartype] = [1,1]
					#print("12", Trios[j].pro_haps[Vartype], gt_prob, gt_fa, gt_mo)
					if gt_fa[0] == i+1 and gt_fa[1] == i+1:
						Trios[j].fa_haps[Vartype] = [1,1]
					else:
						Trios[j].fa_haps[Vartype][0] = 1
					if gt_mo[0] == i+1 and gt_mo[1] == i+1:
						Trios[j].mo_haps[Vartype] = [1,1]
					else:
						Trios[j].mo_haps[Vartype][0] = 1

				elif gt_prob[1] == gt_fa[1] and gt_mo[1] == 0 and gt_prob[1] == i+1 : #paternal transmitted
					if gt_fa[0] == i+1: # transmitted from hom parternal
						Trios[j].fa_haps[Vartype] = [1,1]
					else: # transmitted from het paternal
						Trios[j].fa_haps[Vartype][0] = 1
					Trios[j].pro_haps[Vartype][0] = 1

				#elif gt_fa[1] == 0 and gt_prob[1] == 0:
				#	Trios[j].fa_haps[Vartype][1] = 1

				elif gt_prob[1] == gt_mo[1] and gt_fa[1] == 0 and gt_prob[1] == i+1: #maternal transmitted
					if gt_mo[0] == i+1: # transmitted from hom maternal
						Trios[j].mo_haps[Vartype] = [1,1]
					else: # transmitted from het maternal
						Trios[j].mo_haps[Vartype][0] = i+1
					Trios[j].pro_haps[Vartype][1] = 1

				elif gt_prob[1] == 0: # proband 0/0
					if gt_fa[1] == 1: # father has one hap
						if Trios[j].fa_haps[Vartype][0] == 1:
							Trios[j].fa_haps[Vartype][1] = 0
						elif Trios[j].fa_haps[Vartype][0] == 0:
							Trios[j].fa_haps[Vartype][1] = 1
					if gt_mo[1] == 1: # mother has one hap
						if Trios[j].mo_haps[Vartype][0] == 1:
							Trios[j].mo_haps[Vartype][1] = 0
						elif Trios[j].mo_haps[Vartype][0] == 0:
							Trios[j].mo_haps[Vartype][1] = 1

class Sample:
	def __init__(self, row):
		#self.FamID = row["FamID"]
		#self.sampleID = row["SampleID"]
		#self.Father = row["Paternal"]
		#self.Mother = row["Maternal"]
		#self.Sex = row["Sex"]
		#self.Affected = row["Affected"]
		self.FamID = row[0]
		self.sampleID = row[1]
		self.Father = row[2]
		self.Mother = row[3]
		self.Sex = row[4]
		self.Affected = row[5]
		self.index = row[-1]
	def show(self):
		print(self.FamID, self.sampleID, self.Father, self.Mother, self.Sex, self.Affected)
	def display(self):
		#return "\t".join([self.FamID, self.sampleID, self.Father, self.Mother, self.Sex, self.Affected])
		return list(map(str, [self.FamID, self.sampleID, self.Father, self.Mother, self.Sex, self.Affected]))

class Family:
	def __init__(self, FamID):
		self.FamID = FamID
		self.Father = None
		self.Mother = None
		self.Proband = None
		self.Siblings = []
		self.pro_haps = {}
		self.fa_haps = {}
		self.mo_haps = {}
	def show(self):
		print("FamID:{} Proband:{} Father:{} Mother:{} Siblings:{}".format(
			self.FamID, self.Proband.sampleID, self.Father.sampleID, self.Mother.sampleID, ", ".join(
				[x.sampleID for x in self.Siblings])))

class GTFRecord:
	def __init__(self, Chr, source, Type, start, end, strand, info):
		self.Chr = Chr
		self.source = source
		self.Type = Type
		self.start = start
		self.end = end
		self.strand = strand
		self.info = info

def gtf_info_parser(info):
	res = {}
	for term in info.split(";"):
		if term == "":
			continue
		#print(">",term)
		key,v = term.split()
		v = v.strip('"')
		res[key]=v
	return res

def LoadGeneCode(genecodefil):
	Genes = {}
	Transcripts = {}
	hand = open(genecodefil, 'rt')
	for l in hand:
		if l.startswith("#"):
			continue
		llist = l.strip().split("\t")
		info = gtf_info_parser(llist[8])
		if llist[2] == "gene":
			Genes[info["gene_name"]] = GTFRecord(llist[0], llist[1], llist[2], llist[3], llist[4], llist[6], info)
			Transcripts[info["gene_name"]] = []
		elif llist[2] == "transcript":
			if info["gene_name"] not in Genes:
				Genes[info["gene_name"]] = GTFRecord(llist[0], llist[1], llist[2], llist[3], llist[4], llist[6], info)
				Transcripts[info["gene_name"]] = []
			Transcripts[info["gene_name"]].append(GTFRecord(llist[0], llist[1], llist[2], llist[3], llist[4], llist[6], info))
	return Genes, Transcripts
#
def LoadDeepVar(SPLIST):
	DeepVarFil = open("/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/DeepVariant/SPID2DV.txt", 'rt')
	Dict = {}
	for l in DeepVarFil:
		SPID, FILE = l.strip().split(",")
		if SPID in SPLIST:
			#print(FILE)
			Dict[SPID] = pysam.TabixFile(FILE) #FILE
	return Dict

def GetOptions():
	parser = argparse.ArgumentParser()
	#parser.add_argument("-v", "--vcf", type=str, required=True, help="<Required> VCF file")
	#parser.add_argument("-l", "--list", type=str, required=True, help="<Required> Indv List file")
	#parser.add_argument("-o", "--out", type=str, help="<Required> Output VCF file")
	parser.add_argument("--chr", required=True, type=str, help="<Required> Chromosome")
	parser.add_argument("--af", default=1e-2, type=float, help="<Required> MAF cutoff")
	args = parser.parse_args()
	#if args.out == None:
	#	args.out = "test.out.vcf"
	return args

def main():
	args = GetOptions()
	ins = RecessiveModel(args.af)
	#List = [l.strip() for l in open(args.list)]
	#ins.ComputeSiteAF(args.vcf, List, "EUR", args.out)
	Chr = args.chr
	#GenotypeFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K/VCF/TrioVCF/Genotypes/SPARK30K.TrioSamples.Chr{}.vcf.gz".format(Chr)
	#VEPFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K/VCF/TrioVCF/sites/SPARK30K.TrioSamples.Chr{}.vep.mappability.vcf.gz".format(Chr)
	#VEPFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K/VCF/TrioVCF/sites/Annotated2/SPARK30K.TrioSamples.Chr{}.vep.vcf.gz".format(Chr)
	#AFFil = "/home/local/users/jw/Genetics_Projects/SPARK/spark_genomics/dat/SPARK30K.TrioSamples.Chr{}.eurAF.vcf.gz".format(Chr)
	genecode = "/home/local/users/jw/vep_data/homo_sapiens/GeneCodeV29/CHRs/genecodev29.{}.gtf".format(Chr)
	#GenotypeFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/GenotypesSplitbyChr/GATK4_20190729.chr{}.vcf.gz".format(Chr)
	#VEPFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/SitesSplitbyChr/annotated/GATK4_20190729.chr{}.mappability.vcf.gz".format(Chr)
	GenotypeFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/Filt/SPARK30K.Genotypes.Filt.Chr{}.vcf.gz".format(Chr)
	VEPFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/Filt/SPARK30K.Anno.Filt.Chr{}.vcf.gz".format(Chr)
	#AFFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/recessive/AF/GATK4_20190729.chr{}.eurAF.vcf.gz".format(Chr)
	#ins.Recessive(Chr, GenotypeFil ,VEPFil, AFFil, genecode)
	ins.Recessive(Chr, GenotypeFil ,VEPFil, genecode)

if __name__=='__main__':
	main()


import numpy as np
from rlgym.envs import Match
from rlgym.utils.action_parsers import DiscreteAction
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.vec_env import VecMonitor, VecNormalize, VecCheckNan
from stable_baselines3.ppo import MlpPolicy

from rlgym.utils.obs_builders import AdvancedObs
from rlgym.utils.state_setters import DefaultState
from rlgym.utils.terminal_conditions.common_conditions import TimeoutCondition, NoTouchTimeoutCondition, GoalScoredCondition
from rlgym_tools.sb3_utils import SB3MultipleInstanceEnv
from rlgym.utils.reward_functions.common_rewards.misc_rewards import EventReward
from rlgym.utils.reward_functions.common_rewards.player_ball_rewards import VelocityPlayerToBallReward
from rlgym.utils.reward_functions.common_rewards.ball_goal_rewards import VelocityBallToGoalReward
from rlgym.utils.reward_functions import CombinedReward

if __name__ == '__main__':  # Required for multiprocessing
    frame_skip = 8          # Number of ticks to repeat an action
    half_life_seconds = 5   # Easier to conceptualize, after this many seconds the reward discount is 0.5

    fps = 120 / frame_skip
    gamma = np.exp(np.log(0.5) / (fps * half_life_seconds))  # Quick mafs
    agents_per_match = 2
    num_instances = 1
    target_steps = 1_000_000
    steps = target_steps // (num_instances * agents_per_match) #making sure the experience counts line up properly
    batch_size = target_steps//10 #getting the batch size down to something more manageable - 100k in this case
    training_interval = 25_000_000
    mmr_save_frequency = 50_000_000

    def exit_save(model):
        model.save("models/exit_save")

    def get_match():  # Need to use a function so that each instance can call it and produce their own objects
        return Match(
            team_size=1,
            tick_skip=frame_skip,
            reward_function=CombinedReward(
            (
                VelocityPlayerToBallReward(),
                VelocityBallToGoalReward(),
                EventReward(
                    team_goal=100.0,
                    concede=-100.0,
                    shot=5.0,
                    save=30.0,
                    demo=10.0,
                ),
            ),
            (0.1, 1.0, 1.0)),
            # self_play=True,  in rlgym 1.2 'self_play' is depreciated. Uncomment line if using an earlier version
            terminal_conditions=[TimeoutCondition(fps * 300), NoTouchTimeoutCondition(fps * 45), GoalScoredCondition()],
            obs_builder=AdvancedObs(),  # Not that advanced, good default
            state_setter=DefaultState(),  # Resets to kickoff position
            action_parser=DiscreteAction()  # Discrete > Continuous don't @ me
        )

    env = SB3MultipleInstanceEnv(get_match, num_instances)            # Start 1 instances, waiting 60 seconds between each
    env = VecCheckNan(env)                                # Optional
    env = VecMonitor(env)                                 # Recommended, logs mean reward and ep_len to Tensorboard
    env = VecNormalize(env, norm_obs=False, gamma=gamma)  # Highly recommended, normalizes rewards

    try:
        model = PPO.load(
            "models/exit_save.zip",
            env,
            device="auto",
            custom_objects={"n_envs": env.num_envs}, #automatically adjusts to users changing instance count, may encounter shaping error otherwise
            # If you need to adjust parameters mid training, you can use the below example as a guide
            #custom_objects={"n_envs": env.num_envs, "n_steps": steps, "batch_size": batch_size, "n_epochs": 10, "learning_rate": 5e-5}
        )
        print("Loaded previous exit save.")
    except:
        print("No saved model found, creating new model.")
        from torch.nn import Tanh
        policy_kwargs = dict(
            activation_fn=Tanh,
            net_arch=[512, 512, dict(pi=[256, 256, 256], vf=[256, 256, 256])],
        )

        model = PPO(
            MlpPolicy,
            env,
            n_epochs=10,                 # PPO calls for multiple epochs
            policy_kwargs=policy_kwargs,
            learning_rate=5e-5,          # Around this is fairly common for PPO
            ent_coef=0.01,               # From PPO Atari
            vf_coef=1.,                  # From PPO Atari
            gamma=gamma,                 # Gamma as calculated using half-life
            verbose=3,                   # Print out all the info as we're going
            batch_size=batch_size,             # Batch size as high as possible within reason
            n_steps=steps,                # Number of steps to perform before optimizing network
            tensorboard_log="logs",  # `tensorboard --logdir out/logs` in terminal to see graphs
            device="auto"                # Uses GPU if available
        )

    # Save model every so often
    # Divide by num_envs (number of agents) because callback only increments every time all agents have taken a step
    # This saves to specified folder with a specified name
    callback = CheckpointCallback(round(5_000_000 / env.num_envs), save_path="models", name_prefix="rl_model")

    try:
        mmr_model_target_count = model.num_timesteps + mmr_save_frequency
        while True:
            #may need to reset timesteps when you're running a different number of instances than when you saved the model
            model.learn(training_interval, callback=callback, reset_num_timesteps=False) #can ignore callback if training_interval < callback target
            model.save("models/exit_save")
            if model.num_timesteps >= mmr_model_target_count:
                model.save(f"mmr_models/{model.num_timesteps}")
                mmr_model_target_count += mmr_save_frequency

    except KeyboardInterrupt:
        print("Exiting training")

    print("Saving model")
    exit_save(model)
    print("Save complete")


import numpy as np
import pandas as pd
from pandas import Series, DataFrame

# Let's create a dframe to work with (Highest elevation cities in USA)
dframe = DataFrame({'city': ['Alma', 'Brian Head', 'Fox Park'],
                    'altitude': [3158, 3000, 2762]})

# Now let's say we wanted to add a column for the States, we can do that
# with a mapping.
state_map = {'Alma': 'Colorado', 'Brian Head': 'Utah', 'Fox Park': 'Wyoming'}

# Now we can map that data to our current dframe
dframe['state'] = dframe['city'].map(state_map)

# Mapping is a great way to do element-wise transfomations and other data
# cleaning operations!


import numpy as np

def print_start_message(method_name):
    print('\n\n\n---------- Optimization with {:s} started.\n\n'.format(method_name))

def print_end_message(method_name, time_spent):
    print('\n---------- Training over - {:s}. Took {:d} seconds. \n\n'.format(method_name, np.math.ceil(time_spent)))

def print_progress(i, maxit, val_F, val_f, val_g):
    print('Iter = {:d}/{:d}, F(X) = {:f}, f(X) = {:f}, g(X) = {:f}'.format(i, maxit, val_F, val_f, val_g))


import numpy as np

# Information about the memory layout of the array.

# The flags object can be accessed dictionary-like (as in a.flags['WRITEABLE']), or by using
# lowercased attribute names (as in a.flags.writeable). Short flag names are only supported
# in dictionary access.

a = np.arange(10)
# The data is in a single, C-style contiguous segment
print("Flags Attribute C_CONTIGUOUS : ",a.flags.c_contiguous)
# The data is in a single, Fortran-style contiguous segment
print("Flags Attribute F_CONTIGUOUS : ",a.flags.f_contiguous)
# The array owns the memory it uses or borrows it from another object
print("Flags Attribute OWNDATA : ",a.flags.owndata)
# The data area can be written to. Setting this to False locks the data, making it read-only
print("Flags Attribute WRITEABLE : ",a.flags.writeable)
# The data and all elements are aligned appropriately for the hardware
print("Flags Attribute ALIGNED : ",a.flags.aligned)
# This array is a copy of some other array. When this array is deallocated, the base array
# will be updated with the contents of this array
print("Flags Attribute WRITEBACKIFCOPY  : ",a.flags.writebackifcopy)


# coding=utf-8
import pandas as pd
import numpy as np

df1 = pd.DataFrame({'Note': [111, 222], 'Weekday': ['Mon', 'Tue']}, index=[1, 2])
df2 = pd.DataFrame({'Note': [333, 444], 'Weekday': ['Wed', 'Thu']}, index=[3, 4])
df3 = pd.DataFrame({'Note': [555, 666], 'Weekday': ['Fri', 'Sat']}, index=[5, 6])
df4 = pd.DataFrame({'Note': [777], 'Weekday': ['Sun']}, index=[7])

# concat
df_concat = pd.concat([df1, df2, df3, df4], keys=['df1', 'df2', 'df3', 'df4'])
print("df_concat=\n{}\n".format(df_concat))  # concat函数默认将多个数据纵向串联，结果存在MultiIndex
df_concat_column = pd.concat([df1, df2, df3, df4], axis=1)  # 指定参数“axis=1”以列为主进行串联
print("df_concat_column=\n{}\n".format(df_concat_column))  # 其余数据以“NaN”填充

# append
df_append = df1.append([df2, df3, df4])  # append函数在原数据基础上添加其他数据来进行串联，结果与原数据结构相同
print("df_append=\n{}\n".format(df_append))

# merge
df5 = pd.DataFrame({'Key': ['key1', 'key2', 'key3'], 'A': ['a1', 'a2', 'a3'], 'B': ['b1', 'b2', 'b3']})
df6 = pd.DataFrame({'Key': ['key3', 'key6', 'key9'], 'A': ['a3', 'a6', 'a9'], 'B': ['b3', 'b6', 'b9']})
print("df5=\n{}\ndf6=\n{}\n".format(df5, df6))
merge_df = pd.merge(df5, df6)  # merge函数默认inner方式（内连接，也就是获得键的交集）
merge_inner_key = pd.merge(df5, df6, how='inner', on=['Key'])  # 参数on表示将当作连接键的列名
print("merge_df=\n{}\nmerge_inner=\n{}\n".format(merge_df, merge_inner_key))
merge_left = pd.merge(df5, df6, how='left')  # 左连接
merge_left_key = pd.merge(df5, df6, how='left', on=['Key'])
print("merge_left=\n{}\nmerge_left_key=\n{}\n".format(merge_left, merge_left_key))
merge_right = pd.merge(df5, df6, how='right')  # 右连接
merge_right_key = pd.merge(df5, df6, how='right', on=['Key'])
print("merge_right=\n{}\nmerge_right_key=\n{}\n".format(merge_right, merge_right_key))
merge_outer = pd.merge(df5, df6, how='outer')  # 外连接，也就是获得键的并集
merge_outer_key = pd.merge(df5, df6, how='outer', on=['Key'])
print("merge_outer=\n{}\nmerge_outer_key=\n{}\n".format(merge_outer, merge_outer_key))

# join
df7 = pd.DataFrame({'Key': ['key1', 'key2', 'key3'],
                    'A': ['a1', 'a2', 'a3'],
                    'B': ['b1', 'b2', 'b3']},
                   index=[0, 1, 2])  # 注意索引
df8 = pd.DataFrame({'Key': ['key3', 'key6', 'key9'],
                    'A': ['a3', 'a6', 'a9'],
                    'B': ['b3', 'b6', 'b9']},
                   index=[1, 2, 3])  # join函数可被用于合并多个DataFrame（具有相同的或者类似索引）
print("df7=\n{}\ndf8=\n{}\n".format(df7, df8))
join_left = df7.join(df8, how='left', lsuffix='_self', rsuffix='_other')  # join函数根据索引进行数据合并，默认是左连接
join_right = df7.join(df8, how='right', lsuffix='_self', rsuffix='_other')  # 通过参数lsuffix和rsuffix可以区分结果的列名
print("join_left=\n{}\njoin_right=\n{}\n".format(join_left, join_right))
join_inner = df7.join(df8, how='inner', lsuffix='_self', rsuffix='_other')
join_outer = df7.join(df8, how='outer', lsuffix='_self', rsuffix='_other')
print("join_inner=\n{}\njoin_outer=\n{}\n".format(join_inner, join_outer))

# ### 数据整合
# http://pandas.pydata.org/pandas-docs/stable/merging.html
#
# ### 理解merge和join中的连接方式
# - inner：只返回两张表中都满足on条件的记录
# - outer：返回两张表中的所有记录，对于不满足on条件一端的记录用NaN替换
# - left：返回“表名1”的全部行，对于“表名2”中，不满足on条件的记录用NaN替换
# - right：返回“表名2”的全部行，对于“表名1”中，不满足on条件的记录用NaN替换


#!/usr/bin/env python
# coding: utf-8

#
#
# # Earth Analytics Education - EA  Python Course Spring 2021

# ## Important  - Assignment Guidelines
#
# 1. Before you submit your assignment to GitHub, make sure to run the entire notebook with a fresh kernel. To do this first, **restart the kernel** (in the menubar, select Kernel$\rightarrow$Restart & Run All)
# 2. Always replace the `raise NotImplementedError()` code with your code that addresses the activity challenge. If you don't replace that code, your notebook will not run.
#
# ```
# # YOUR CODE HERE
# raise NotImplementedError()
# ```
#
# 3. Any open ended questions will have a "YOUR ANSWER HERE" within a markdown cell. Replace that text with your answer also formatted using Markdown.
# 4. **DO NOT RENAME THIS NOTEBOOK File!** If the file name changes, the autograder will not grade your assignment properly.
# 6. When you create a figure, comment out `plt.show()` to ensure the autograder can grade your plots. For figure cells, DO NOT DELETE the code that says `DO NOT REMOVE LINE BELOW`.
#
# ```
# ### DO NOT REMOVE LINE BELOW ###
# student_plot1_ax = nb.convert_axes(plt)
# ```
#
# * Only include the package imports, code, and outputs that are required to run your homework assignment.
# * Be sure that your code can be run on any operating system. This means that:
#    1. the data should be downloaded in the notebook to ensure it's reproducible
#    2. all paths should be created dynamically using the `os.path.join`
#
# ## Follow to PEP 8 Syntax Guidelines & Documentation
#
# * Run the `autopep8` tool on all cells prior to submitting (HINT: hit shift + the tool to run it on all cells at once!
# * Use clear and expressive names for variables.
# * Organize your code to support readability.
# * Check for code line length
# * Use comments and white space sparingly where it is needed
# * Make sure all python imports are at the top of your notebook and follow PEP 8 order conventions
# * Spell check your Notebook before submitting it.
#
# For all of the plots below, be sure to do the following:
#
# * Make sure each plot has a clear TITLE and, where appropriate, label the x and y axes. Be sure to include UNITS in your labels.
#

# ### Add Your Name Below
# **Lauren Kremer**

#

# ---

# # Week 04 and 05 Homework - Automate NDVI Workflow
#
# For this assignment, you will write code to generate a plot of the mean normalized difference vegetation index (NDVI) for two different sites in the United States across one year of data:
#
# * San Joaquin Experimental Range (SJER) in Southern California, United States
# * Harvard Forest (HARV) in the Northeastern United States
#
# The data that you will use for this week is available from **earthpy** using the following download:
#
# `et.data.get_data('ndvi-automation')`
#
# ## Assignment Goals
#
# Your goal in this assignment is to create the most efficient and concise workflow that you can that allows for:
#
# 1. The code to scale if you added new sites or more time periods to the analysis.
# 2. Someone else to understand your workflow.
# 3. The LEAST and most efficient (i.e. runs fast, minimize repetition) amount of code that completes the task.
#
# ### HINTS
#
# * Remove values outside of the landsat valid range of values as specified in the metadata, as needed.
# * Keep any output files SEPARATE FROM input files. Outputs should be created in an outputs directory that is created in the code (if needed) and/or tested for.
# * Use the functions that we demonstrated during class to make your workflow more efficient.
# * BONUS - if you  chose - you can export your data as a csv file. You will get bonus points for doing this.
#
#
# ## Assignment Requirements
#
# Your submission to the GitHub repository should include:
# * This Jupyter Notebook file (.ipynb) with:
#     * The code to create a plot of mean NDVI across a year for  2 NEON Field Sites:
#         * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object
#     * The **data should be cleaned to remove the influence of clouds**. See the [earthdatascience website for an example of what your plot might look like with and without removal of clouds](https://www.earthdatascience.org/courses/earth-analytics-python/create-efficient-data-workflows/).
# * BONUS: Create one output `.csv` file that has 3 columns - NDVI, Date and Site Name - with values for SJER and HARV.
#
# Your notebook should:
# * Have *at least* 2 well documented and well named functions with docstrings.
# * Include a Markdown cell at the top of the notebook that outlines the overall workflow using pseudocode (i.e. plain language, not code)
# * Include additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Replace this cell with your pseudocode  for this workflow
#
# If you happen to be a diagram person a diagram is ok too
#
#

# 1. Open each "sites" folder
#     a. open the landsat-crop folder
#         -generate a list of files for ndvi calcs (bands 4 and 5)
#         -access qa files from same folder for cloud mask
#     b. make a list of paths for crop .shp files
#
# 2.  With bands 4 and 5 for each image, generate a function that:
#     a. mask invalid values (1-10000)
#     b. calculate NDVI
#     c. clip to boundary
#     d. mask clouds using qa layer
#     e. calculate mean NDVI for each image and generate a
#     dataframe with image date and site colums.
#
# 3. apply NDVI function
#     a. generate dateframes from band and crop shape lists in a loop
#     b. concatonate dataframe from each image into one plottable df
#
# 4. Plot ndvi values as a timeseries
#

# In[1]:

# Autograding imports - do not modify this cell
import matplotcheck.autograde as ag
import matplotcheck.notebook as nb
import matplotcheck.timeseries as ts
from datetime import datetime

# In[2]:

# Import needed packages in PEP 8 order
# and no unused imports listed (10 points total)

import os
from glob import glob

import matplotlib.pyplot as plt
import numpy as np
import geopandas as gpd
import rioxarray as rxr
import xarray as xr
import earthpy as et
import pandas as pd

# Download data
et.data.get_data('ndvi-automation')

# Set working directory
os.chdir(os.path.join(et.io.HOME,
                      "earth-analytics",
                      "data"))

# In[3]:

# DO NOT MODIFY THIS CELL
# Tests that the working directory is set to earth-analytics/data

path = os.path.normpath(os.getcwd())
student_wd_parts = path.split(os.sep)

if student_wd_parts[-2:] == ['earth-analytics', 'data']:
    print("\u2705 Great - it looks like your working directory is set correctly to ~/earth-analytics/data")
else:
    print("\u274C Oops, the autograder will not run unless your working directory is set to earth-analytics/data")

# In[4]:

# The NDVI automation download provides 30m resolution Landsat data and shapefiles for
# two study sites, one in San Joaquin, CA and the other in Harvard Forest, MA. The two sites
# differ in vegetation coverage density and type, making them ideal for reviewing potential
# differences in NDVvi time series analysis. Landsat data provided includes images for
# approximately 23 days between Jan 12, 2017 to December 30, 2017 and for each image,
# bands 1-5 and a quality assessment band are included in the download.

# # Figure 1: Plot 1 - Mean NDVI For Each Site Across the Year (50 points)
#
# Create a plot of the mean normalized difference vegetation index (NDVI) for the two different sites in the United States across the year:
#
# * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object.
# * Each site should be identified with a different color in the plot and legend.
# * The final plot **data should be cleaned to remove the influence of clouds**.
# * Be sure to include appropriate title and axes labels.
#
# Add additional cells as needed for processing data (e.g. defining functions, etc), but be sure to:
# * follow the instructions in the code cells that have been provided to ensure that you are able to use the sanity check tests that are provided.
# * include only the plot code in the cell identified for the final plot code below

# ## Task 1:
#
# In the cell below, create a single dataframe containing MEAN NDVI, the site name,
# and the date of the data for the HARV site
# scene `HARV/landsat-crop/LC080130302017031701T1-SC20181023151837`.  The column names for the  final
# DataFrame should be`mean_ndvi`, and `site`, and the data should be **indexed on the date**.
#
# Use the functions that we reviewed in class (or create your own versions of them) to implement your code
#
# ### In the Cell below Place  All Functions Needed to Run this Notebook (20 points)

# In[5]:

### DO NOT REMOVE THIS LINE OR EDIT / MOVE THIS CELL ###
start_time = datetime.now()

# In[6]:

# In this cell place all of the functions needed to run your notebook
# You will be graded here on function application, docstrings, efficiency so ensure
# All functions are placed here!

# Function 1. Generate a crop_boundary for each site

def open_boundary(site_path):
    """generate a list of boundary shapefiles for clipping landsat data to study area
    Parameters
    -----------
    vector_path : a path to the directory containing desire shapefile

    Returns
    -----------
    gpd : a geopandas geodataframe
    """
    # Open crop boundary
    vector_dir = os.path.join(site_path, "vector")
    site_name = os.path.basename(os.path.normpath(site_path))
    site_boundary_path = os.path.join(vector_dir,  site_name + "-crop.shp")
    crop_bound = gpd.read_file(site_boundary_path)
    crop_bound['Site'] = site_name
    return crop_bound

# Function 2. Generate a list of landsat 8 image folders for each site

def build_image_list(landsat_dir):
    """generates a list of landsat bands 4 and 5 from folder of images
    Parameters
    -----------
    landsat_dir : a path to the folder/directory of image subfolders which
    contain landsat bands

    Returns
    -----------
    list : a list of .tif filepaths that can be used to calculate NDVI
    """

    image_list = []

    image_paths = sorted(glob(os.path.join(landsat_dir, "LC08*")))
    image_list.append(image_paths)
    image_list = [i for b in map(lambda x:[x] if not isinstance(x, list)
                                 else x, image_list) for i in b] # unlists nested lists
    return image_list

# Function 3.
def mask_crop_ndvi(image_folder, vals, crop_extent, valid_range = None):

    """Open bands 4 and 5 in each image folder, mask according to valid_range of pixel
    values, then stores them in a two band list used to calculate NDVI at each pixel.
    Then masks clouds using the qa file in the image folder and clips them from a
    .shp file indicated by the function parameter (crop_extent). Then
    generates dataframe from image mean NDVI with date and site name.

    Parameters
    -----------
    image_folder : list of Landsat 8 folders containing image bands
    vals: A list of values needed to create the cloud mask
    crop_extent: geopandas GeoDataFrame
        A geopandas dataframe to be used to crop the raster data
    valid_range : tuple (optional)
        A tuple of min and max range of values for the data. Default = None

    Returns
    -----------
    ndvi_df : Pandas dataframe
        a dataframe containing cropped image mean NDVI values indexed by date
    """

    band_path = sorted(glob(os.path.join(image_folder, "*band*[4-5].tif")))
    opened_bands = []

    for i in band_path:
        landsat_45 = rxr.open_rasterio(i, masked=True).rio.clip(crop_extent.geometry,
                                                            from_disk=True).squeeze()
        if valid_range:
            mask = ((landsat_45 < valid_range[0]) | (landsat_45 > valid_range[1]))
            band = landsat_45.where(~xr.where(mask, True, False))
            opened_bands.append(band)

    img_ndvi = (opened_bands[1] - opened_bands[0]) / (opened_bands[1] + opened_bands[0])

    # # then open pixel_qa
    qa_paths = glob(os.path.join(image_folder, "*qa*"))
    path_str = ' '.join([str(elem) for elem in qa_paths])
    qa_open = rxr.open_rasterio(path_str).rio.clip(crop_extent.geometry,
                                                              from_disk=True).squeeze()
    ndvi_mask = img_ndvi.where(~qa_open.isin(vals))

    # create the dataframe
    ndvi_df = pd.DataFrame()
    ndvi_df['mean_ndvi'] = pd.Series({"mean": ndvi_mask.mean()}, dtype=float)
    ndvi_df['site'] = image_folder[22:26]
    ndvi_df['date'] = pd.to_datetime(image_folder[50:58], format = '%Y%m%d')
    ndvi_df.set_index('date', inplace=True)
    return ndvi_df

# In[7]:

# For optimal processing, only bands needed for NDVI calculations and cloud masking (qa)
# were opened. Bands were clipped upon opening to decrease file size to only geographic
# area of interest, and the cloudmask was applied only to the the NDVI array to reduce
# processing time over applying a mask to each band used for the calculation. The output
# of the above function only returns a dataframe of required data, rather than any arrays
# used for calculations within the dataframe.

# In[8]:

# Create dataframe of mean NDVI in this cell using the functions created above
# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Call the dataframe at the end of the cell so the tests run on it!
# Be sure that the date column is an index of type date
# HINT: the time series lessons may help you remember how to do this!

# Indicate path to site directories
path = os.path.join("ndvi-automation", "sites")

# Get a list of both site directories
sites = sorted(glob(path + "/*/"))

# build list (build_image_list)of landsat image folders containing bands:
sitelists = []
for i in sites:
    site_list = build_image_list(landsat_dir = os.path.join(i, "landsat-crop"))
    sitelists.append(site_list)
    #sitelists = [i for b in map(lambda x:[x] if not isinstance(x, list) else x, sitelists) for i in b]
sitelists[0][4] #identify one image folder specific to Task 1.

# Open shapefiles using open_boundary function
bounds = []
for i in sites:
    boundy = open_boundary(i)
    bounds.append(boundy)

bounds[0] # identify harv geodf for Task 1.

# Identify variables for mask_cloud function:

# Cloud no data vals for Landsat 8 -
vals = [328, 392, 840, 904, 1350, 352, 368, 416,
        432, 480, 864, 880, 928, 944, 992, 480, 992]

# apply ndvi calculation and masking function for specific image
mask_crop_ndvi(image_folder = sitelists[0][4], vals = vals, crop_extent = bounds[0],
               valid_range = (0, 10000))

# In[9]:

# This cell  is testing your data output above

student_ndvi_ts_single_site = _

single_scene_points = 0

# Ensure the data is stored in a dataframe.
if isinstance(student_ndvi_ts_single_site, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    single_scene_points += 1
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Ensure that the date column is the index
if isinstance(student_ndvi_ts_single_site.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    single_scene_points += 2
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_ts_single_site.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    single_scene_points += 2
else:
    print('\u274C The data in your date column is not datetime.')

# Ensure the site name is correct
if student_ndvi_ts_single_site.site.values[0] == 'HARV':
    print('\u2705 You have the correct site name!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct site name.')

if np.allclose(0.281131628228094, student_ndvi_ts_single_site.mean_ndvi.values[0]):
    print('\u2705 You have the correct mean NDVI value!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct mean ndvi value.')

print("\n \u27A1 You received {} out of 15 points for creating a dataframe.".format(
    single_scene_points))
single_scene_points

# ## Task 2:
#
# In the cell below, process all of the landsat scenes. Create a DataFrame that contains the following
# information for each scene
#
#
# |   | index  | site  | mean_ndvi  |
# |---|---|---|---|
# | Date  |   |   |   |
# |  2017-01-07  | 0  | SJER  | .4  |
#
# Be sure to call your dataframe at the end of the cell to ensure autograding works.
# HINT: FOR THIS STEP, leave any rows containing missing values (`NAN`).

# In[10]:

# Create dataframe of NDVI including the cleaning data to deal with clouds

# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Don't forget to set date as the index and make the values of type datetime

# apply ndvi calc and masking function using bounds list as the crop-extent parameter
#in mask_crop_ndvi
indexed_dfs = []

for i, j  in zip(sitelists, bounds):
    for k in i:
        output_df = mask_crop_ndvi(image_folder = k, vals = vals, crop_extent = j,
                   valid_range = (0, 10000))
        indexed_dfs.append(output_df) #add it to a new list sorted in order

mean_masked_ndvi = pd.concat(indexed_dfs) #combine dfs in list to new dataframe
mean_masked_ndvi # call final object

# In[11]:

# Using a parallel loop to call the appropriate boundary shapefile in each iteration
# cloud_mask function may not be ideal for processing time (loops costly in terms
# of processing time?), but the script is efficiently short, easy to follow and produces
# the entire dataframe with one command (at least before concatonation of the dataframe),
# rather than having to run the mask function separately for each site/shapefile.

# In[12]:

# Last sanity check before creating your plot (10 points)

# Ensure that you call your dataframe at the bottom of the cell above
# and that it has columns called: mean_ndvi and site

# Ensure the data is stored in a dataframe.
student_ndvi_df = _

df_points = 0

if isinstance(student_ndvi_df, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    df_points +=2
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Check that dataframe contains the appropriate number of NAN values
if student_ndvi_df.mean_ndvi.isna().sum() == 15:
    print('\u2705 Correct number of masked data values!')
    df_points +=2
else:
    print('\u274C The amount of null data in your dataframe is incorrect.')

# Ensure that the date column is the index
if isinstance(student_ndvi_df.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    df_points +=3
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_df.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    df_points +=3
else:
    print('\u274C The data in your date column is not datetime.')

# Output for timer, # DO NOT MODIFY
end_time = datetime.now()
total_time = end_time - start_time
print(
    "Your total run time for processing the data was {0}.".format(total_time))

print("\n \u27A1 You received {} out of 10 points for creating a dataframe.".format(
    df_points))

df_points

# In[13]:

# Add only the plot code to this cell

# This is the final figure of mean NDVI
# for both sites across the year
# with data cleaned to deal with clouds

# subset dataframe to exclude "NA" values
plot_data = mean_masked_ndvi[mean_masked_ndvi['mean_ndvi'].notna()]

colorPalette = {'HARV': 'magenta',
                'SJER': 'navy'}

# # Create figure and plot space
fig, ax = plt.subplots(figsize=(10, 6))

for key, data in plot_data.groupby('site'):
    data.plot(use_index=True,
              y='mean_ndvi',
              ax=ax,
              marker='o',
              label=key,
              color = colorPalette)

### DO NOT REMOVE LINES BELOW ###
final_masked_solution = nb.convert_axes(plt, which_axes="current")

# In[14]:

# Ignore this cell for the autograding tests

# In[15]:

# Ignore this cell for the autograding tests

# # Question 1 (10 points)
#
# Imagine that you are planning NEON’s upcoming flight season to capture remote sensing data in these locations and want to ensure that you fly the area when the vegetation is the most green.
#
# When would you recommend the flights take place for each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# 1. I would opt to fly at a time near the peak mean NDVI, likely early April for San Joaquin and near the end of June for Harvard Forest.  However, I would likely pick a period from several years worth of data.  Using Landsat 8, we may be able to determine an ideal period from up to eight years of data.

# # Question 2 (10 points)
#
# How could you modify your workflow to look at vegetation changes over time in each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# 2. My current workflow would support that pretty easily. To look at an individual site, I could skip the part of the loops that incorporated 'sites' which was a means to access multiple sites. The current functions review all image folders in the supplied path, so we could extend the time period reviewed simply by adding more image folders. If I wanted to review NDVI images rather than compare changes in a dataframe, I could remove the dataframe building chunk from my mask_cloud function and the function would instead return a list of arrays that I could plot from.

# # Do not edit this cell! (10 points)
#
# The notebook includes:
# * additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Do not edit this cell! (20 points)
#
# The notebook will also be checked for overall clean code requirements as specified at the **top** of this notebook. Some of these requirements include (review the top cells for more specifics):
#
# * Notebook begins at cell [1] and runs on any machine in its entirety.
# * PEP 8 format is applied throughout (including lengths of comment and code lines).
# * No additional code or imports in the notebook that is not needed for the workflow.
# * Notebook is fully reproducible. This means:
#    * reproducible paths using the os module.
#    * data downloaded using code in the notebook.
#    * all imports at top of notebook.

# ## BONUS - Export a  .CSV File to Share (10 points possible)
#
# This is optional - if you export a **.csv** file with the columns specified above: Site, Date and NDVI Value you can get an additional 10 points.
#
# * FULL CREDIT: File exists in csv format and contains the columns specified.
# We will check your github repo for this file!
#

# In[16]:

#csvfile= os.path.join(et.io.HOME,
#                      "earth-analytics",
#                      "ea-2021-04-ndvi-automation-streamfireflies",
#                      "mean_masked_ndvi.csv")
#mean_masked_ndvi.to_csv(csvfile)

# Was able to generate a .csv to be pushed to my repo, but the CI didn't like this.
# Is it because the file isn't in the main (earthlab) repo?

# In[ ]:


from problem1 import *
import numpy as np
import sys

'''
    Unit test 1:
    This file includes unit tests for problem1.py.
    You could test the correctness of your code by typing `nosetests -v test1.py` in the terminal.
'''

#-------------------------------------------------------------------------
def test_python_version():
    ''' ----------- Problem 1 (30 points in total)---------------------'''
    assert sys.version_info[0]==3 # require python 3 (instead of python 2)

#-------------------------------------------------------------------------
def test_R_play():
    '''(2 points) Player Random play()'''
    p = PlayerRandom()
    s=np.array([[ 0, 1, 1],
                [ 1, 0,-1],
                [ 1, 1, 0]])
    count=np.zeros(3)
    for _ in range(100):
        r,c = p.play(s)
        assert s[r,c]==0
        assert r==c
        assert r>-1 and r<3
        count[c]+=1
    assert count[0]>20
    assert count[1]>20
    assert count[2]>20

    s=np.array([[ 1, 1, 0],
                [ 1, 0,-1],
                [ 0, 1, 1]])

    for _ in range(100):
        r,c = p.play(s)
        assert s[r,c]==0
        assert r==2-c
        assert r>-1 and r<3

#-------------------------------------------------------------------------
def test_T_play_x():
    '''(2 points) TicTacToe play_x()'''
    g = TicTacToe()
    g.play_x(0,0)
    assert np.allclose(g.s[0,0],1)
    assert np.allclose(g.s.sum(),1)
    g.play_x(2,1)
    assert np.allclose(g.s[2,1],1)
    assert np.allclose(g.s.sum(),2)

#-------------------------------------------------------------------------
def test_T_play_o():
    '''(2 points) TicTacToe play_o()'''
    g = TicTacToe()
    g.play_o(0,0)
    assert np.allclose(g.s[0,0],-1)
    assert np.allclose(g.s.sum(),-1)
    g.play_o(2,2)
    assert np.allclose(g.s[2,2],-1)
    assert np.allclose(g.s.sum(),-2)

#-------------------------------------------------------------------------
def test_T_check():
    '''(3 points) TicTacToe check()'''
    g = TicTacToe()
    e = g.check(g.s)
    assert e is None
    g.play_x(0,0)
    g.play_x(0,1)
    g.play_x(0,2)
    e = g.check(g.s)
    assert e == 1

    g = TicTacToe()
    g.play_o(0,0)
    g.play_o(1,0)
    g.play_o(2,0)
    e = g.check(g.s)
    assert e == -1

    g = TicTacToe()
    g.play_o(0,2)
    g.play_o(1,1)
    g.play_o(2,0)
    e = g.check(g.s)
    assert e == -1

    g = TicTacToe()
    g.play_o(2,2)
    g.play_o(1,1)
    g.play_o(0,0)
    e = g.check(g.s)
    assert e == -1

    g = TicTacToe()
    g.s=np.array([[-1, 1,-1],
                  [-1, 1,-1],
                  [ 1,-1, 1]])
    e = g.check(g.s)
    assert e == 0

    g.s=np.array([[-1, 0,-1],
                  [-1, 1,-1],
                  [ 1,-1, 1]])
    e = g.check(g.s)
    assert e is None

#-------------------------------------------------------------------------
def test_T_game():
    '''(5 points) TicTacToe game()'''

    # if the game has already ended
    g = TicTacToe()
    p1 = PlayerRandom()
    g.s=np.array([[ 0, 1, 1],
                  [-1,-1,-1],
                  [ 1,-1, 1]])
    e = g.game(p1,p1)
    assert e==-1

    p1 = PlayerRandom()
    p2 = PlayerRandom()
    w =0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0,-1, 1],
                      [-1, 1, 0],
                      [-1, 1,-1]])
        e = g.game(p1,p2)
        w+=e
    print(w)
    assert w<0
    assert w<-30
    assert w>-70

    class test1:
        def play(self,s):
            assert s[1,1] == 1
            assert s[2,2] ==-1
            r,c=np.where(s==0)
            return r[0],c[0]
    class test2:
        def play(self,s):
            assert s[1,1] ==-1
            assert s[2,2] == 1
            r,c=np.where(s==0)
            return r[0],c[0]

    g = TicTacToe()

    p1 = test1()
    p2 = test2()

    g.s=np.array([[ 0, 0, 0],
                  [ 0, 1, 0],
                  [ 0, 0,-1]])

    e = g.game(p1,p2)

    class test3:
        def play(self,s):
            r,c=np.where(s==0)
            return r[0],c[0]

    g = TicTacToe()

    p1 = test3()
    e = g.game(p1,p1)
    print(g.s)
    s=np.array([[ 1,-1, 1],
                [-1, 1,-1],
                [ 1, 0, 0]])
    assert np.allclose(g.s, s)
    assert e==1

#-------------------------------------------------------------------------
def test_M_update_v():
    '''(2 points) Player MiniMax update_v()'''
    p = PlayerMiniMax()
    s=np.array([[ 1, 0, 0],
                [ 0, 1, 0],
                [ 0,-1, 1]])
    p.update_v(s,1) # won the game
    e = p.v[str(s)]
    assert  e== 1
    assert len(p.v.keys())== 1

#-------------------------------------------------------------------------
def test_M_update_p():
    '''(2 points) Player MiniMax update_p()'''
    p = PlayerMiniMax()
    s=np.array([[ 0, 0, 0],
                [ 0, 1, 0],
                [ 0, 1,-1]])
    p.update_p(s,0,1)
    r,c = p.p[str(s)]
    assert  r== 0
    assert  c== 1
    assert len(p.p.keys())== 1

#-------------------------------------------------------------------------
def test_M_compute_v():
    '''(5 points) Player MiniMax compute_v()'''
    p = PlayerMiniMax()
    s=np.array([[ 1, 0, 0],
                [ 0, 1, 0],
                [ 0,-1, 1]])
    v=p.compute_v(s) # won the game
    assert  v== 1

    e = p.v[str(s)]
    assert  e== 1

    assert len(p.v.keys())== 1

    s=np.array([[ 0, 0, 0],
                [-1,-1,-1],
                [ 0, 1, 0]])
    v=p.compute_v(s)
    assert  v==-1

    s=np.array([[-1, 1,-1],
                [-1, 1,-1],
                [ 1,-1, 1]])
    v=p.compute_v(s)
    assert  v==0

    s=np.array([[-1, 1,-1],
                [-1, 1,-1],
                [ 0,-1, 1]])
    v=p.compute_v(s)
    assert  v==0

    s=np.array([[-1,-1, 1],
                [-1, 1,-1],
                [ 0,-1, 1]])
    v=p.compute_v(s)
    assert  v==1

    s=np.array([[ 0, 1,-1],
                [-1,-1, 1],
                [ 0,-1, 1]])
    v=p.compute_v(s)
    assert v==0

    s=np.array([[ 0, 1, 1],
                [-1, 1,-1],
                [ 1,-1, 0]])
    v=p.compute_v(s)
    assert v==1

    s=np.array([[ 0, 0, 1],
                [-1, 1, 0],
                [-1, 0, 0]])
    v=p.compute_v(s)
    assert v==1

#-------------------------------------------------------------------------
def test_M_play():
    '''(2 points) Player MiniMax play()'''
    p = PlayerMiniMax()
    s=np.array([[-1, 1,-1],
                [-1, 1,-1],
                [ 0,-1, 1]])
    r, c = p.play(s)
    assert np.allclose(s,[[-1, 1,-1], [-1, 1,-1], [ 0,-1, 1]])
    assert r==2
    assert c==0

    s=np.array([[-1,-1, 1],
                [-1, 1,-1],
                [ 0,-1, 1]])
    r, c = p.play(s)
    assert r==2
    assert c==0

    p = PlayerMiniMax()
    s=np.array([[ 0,-1, 1],
                [-1, 1,-1],
                [ 0,-1,-1]])
    r, c = p.play(s)
    assert r==2
    assert c==0

    p = PlayerMiniMax()
    s=np.array([[ 0, 1,-1],
                [-1,-1, 1],
                [ 0,-1, 1]])
    r, c = p.play(s)
    assert r==2
    assert c==0

    p = PlayerMiniMax()
    s=np.array([[ 0, 1, 1],
                [-1, 1,-1],
                [-1,-1, 1]])
    r, c = p.play(s)
    assert r==0
    assert c==0

    p = PlayerMiniMax()
    s=np.array([[ 0,-1, 1],
                [-1, 1,-1],
                [-1, 1, 0]])
    r, c = p.play(s)
    assert r==0
    assert c==0

#-------------------------------------------------------------------------
def test_players():
    '''(2 points) random vs Minimax'''

    p1 = PlayerMiniMax()
    p2 = PlayerRandom()
    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0,-1, 1],
                      [-1, 1,-1],
                      [ 0,-1,-1]])
        e = g.game(p1,p2)
        w += e
    assert w==100

    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0,-1, 1],
                      [-1, 1,-1],
                      [-1, 1, 0]])
        e = g.game(p1,p2)
        w += e
    assert w==0

    w=0
    p1 = PlayerMiniMax()
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 1],
                      [ 0,-1, 0],
                      [ 1,-1, 0]])
        e = g.game(p1,p2)
        w += e
    assert np.abs(w-87)<10

#-------------------------------------------------------------------------
def test_players2():
    '''(3 points) Minimax vs Minimax'''
    # NOTE: this test can usually finish within 20 seconds.
    # if your code is very slow (say using more than 1 minute), you may want to check the dictionaries (self.v and self.p) implementation.

    p = PlayerMiniMax()
    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 1],
                      [ 0,-1, 0],
                      [ 1,-1, 0]])
        e = g.game(p,p)
        w += e
    assert w==0

    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 0],
                      [ 0,-1, 0],
                      [ 1, 0, 0]])
        e = g.game(p,p)
        w += e
    assert w==0

    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 0],
                      [ 0, 0, 0],
                      [ 1,-1, 0]])
        e = g.game(p,p)
        w += e
    assert w==100

    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 0],
                      [ 0, 1, 0],
                      [ 0,-1, 0]])
        e = g.game(p,p)
        w += e
    assert w==100

    w=0
    for i in range(100):
        g = TicTacToe()
        g.s=np.array([[ 0, 0, 0],
                      [ 0, 1, 0],
                      [-1, 0, 0]])
        e = g.game(p,p)
        w += e
    assert w==0

    w=0
    for i in range(100):
        g = TicTacToe()
        e = g.game(p,p)
        w += e
    assert w==0


import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure(figsize=(4.25,8*.55))
ax = fig.add_axes([0,0,1,1], xlim=[0,11], ylim=[0.5,8.5], frameon=False,
                      xticks=[], yticks=[]) #, aspect=1)

X = np.linspace(1,10,10)
Y = np.zeros(len(X))
y = 8

# Marker edge color
# ----------------------------------------------------------------------------
C = ["C%d" % i for i in range(10)]
plt.scatter(X, y+Y, s=200, facecolor="white", edgecolor=C, linewidth=1.5)
for x,c in zip(X,C):
    plt.text(x, y-0.25, '"%s"' % c,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Marker edge color",
         size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "mec / ec", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Marker face color
# ----------------------------------------------------------------------------
C = ["C%d" % i for i in range(10)]
plt.scatter(X, y+Y, s=200, facecolor=C, edgecolor="None")
for x,c in zip(X,C):
    plt.text(x, y-0.25, '"%s"' % c,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Marker face color",
         size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "mfc / fc", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Marker edge width
# ----------------------------------------------------------------------------
LW = (1+np.arange(10)/2)
plt.scatter(X, y+Y, s=100, facecolor="white", edgecolor="black", linewidth=LW)
for x,lw in zip(X,LW):
    plt.text(x, y-0.25,  "%.1f" % lw,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Marker edge width",
         size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "mew / lw", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Marker edge width
# ----------------------------------------------------------------------------
S = (1+np.arange(10))*25
plt.scatter(X, y+Y, s=S, facecolor="black", edgecolor="None")
for x,s in zip(X,S):
    plt.text(x, y-0.25, '%d' % s,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Marker size",
         size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "ms / s", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

X = np.linspace(1,10,12)

# Filled markers
# -----------------------------------------------------------------------------
M = [".", "o", "s", "P", "X", "*", "p", "D", "<", ">", "^", "v"]
for x, marker in zip(X,M):
    plt.scatter(x, y, s=256, color="black", marker="s",fc=".9", ec="none")
    plt.scatter(x, y, s=100, color="black", marker=marker,
                fc="white", ec="black",linewidth=0.75)
    plt.text(x, y-0.25, '"%s"' % marker,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Filled markers", size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "marker", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Unfilled markers
# -----------------------------------------------------------------------------
M = ["1", "2", "3", "4", "+", "x", "|", "_", 4, 5, 6, 7]
for x, marker in zip(X,M):
    if isinstance(marker,str): text = '"%s"' % marker
    else:                      text = '%s' % marker

    plt.scatter(x, y, s=256, color="black", marker="s",fc=".9", ec="none")
    plt.scatter(x, y, s=100, color="black", marker=marker,
                fc="none", ec="black",linewidth=0.75)
    plt.text(x, y-0.25, text,
             size="x-small", ha="center", va="top", family="monospace")
plt.text(X[0]-0.25, y+0.25, "Unfilled markers", size="small", ha="left", va="baseline")
plt.text(X[-1]+0.25, y+0.25, "marker", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Unicode markers
# -----------------------------------------------------------------------------
M = ["♠","♣","♥","♦", "→","←","↑","↓", "◐","◑","◒","◓"]
for x, marker in zip(X,M):
    ax.scatter(x, y, s=256, color="black", marker="s",fc=".9", ec="none")
    ax.scatter(x, y, s=100, color="black", marker="$"+marker+"$",
                fc="black", ec="none", linewidth=0.5)
    ax.text(x, y-0.25, '"\$%s\$"' % marker,
             size="x-small", ha="center", va="top", family="monospace")
ax.text(X[0]-0.25, y+0.25, "Unicode markers", size="small", ha="left", va="baseline")
ax.text(X[-1]+0.25, y+0.25, "marker", color="blue",
         size="small", ha="right", va="baseline", family="monospace")
y -= 1

# Spacing
# -----------------------------------------------------------------------------
n_segment = 4
width = 9
segment_width = 0.75*(width/n_segment)
segment_pad = (width - n_segment*segment_width)/(n_segment-1)
X0 =  1+np.arange(n_segment)*(segment_width+segment_pad)
marks = [ 10, [0,-1], (25, 5), [0,25,-1] ]

for x0, mark in zip(X0,marks):
    X = np.linspace(x0, x0+segment_width, 50)
    Y = y*np.ones(len(X))
    ax.plot(X, Y, linewidth=1, color="black",
            marker=".", mfc="white", mec="black", mew="1", markevery=mark)

    ax.text((X[0]+X[-1])/2, y-0.1, '%s' % str(mark),
            size="x-small", ha="center", va="top")

ax.text(1-0.25, y+0.25, "Marker spacing", size="small", ha="left", va="baseline")
ax.text(X[-1]+0.25, y+0.25, "markevery", color="blue",
         size="small", ha="right", va="baseline", family="monospace")

plt.savefig("reference-markers.pdf", dpi=600)
plt.show()


# 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