import os
import pkgutil
import re
import shutil
import stat
from string import Template
import subprocess
from typing import List
import warnings
from hepi.input import Input, Order, is_gluino, is_slepton, is_squark, is_weakino
from hepi.results import Result
from hepi.run import RunParam, Runner
from hepi.util import namehash
from smpl import debug
from uncertainties import ufloat



class ProspinoResult(Result):
    """
    Prospino Result Data.
    """

class ProspinoRunner(Runner):
	# TODO add leptoqaurk, stop and sbottom cases
	weakino_map = {
		1000022 : 1,
		1000023 : 2,
		1000025 : 3,
		1000035 : 4,
		1000024 : 5,
		1000037 : 6,
		-1000024 : 7,
		-1000037 : 8,
	}
	squark_map = {
		1000005 : -5,
		1000004 : -4,
		1000004 : -4,
		1000003 : -3,
		1000001 : -2,
		1000002 : -1,
		2000002 : 1,
		2000001 : 2,
		2000003 : 3,
		2000004 : 4,
		2000005 : 5,
	}
	slepton_map = {
		( 1000011, 1000011) :  1,
		( 2000011, 2000011) :  2,
		( 1000012, 1000012) :  3,
		(-1000011, 1000012) :  4,
		( 1000011, 1000012) :  5,
		( 1000015, 1000015) :  6,
		( 2000015, 2000015) :  7,
		( 1000015, 2000015) :  8,
		( 1000016, 1000016) :  9,
		(-1000015, 1000016) : 10,
		( 1000015, 1000016) : 11,
		(-2000015, 1000016) : 12,
		( 2000015, 1000016) : 13,
	}

	def _get_ps_proc(self,p:Input):
		if p.has_squark() and p.has_weakino():
			sq = p.particle1 if is_squark(p.particle1) else p.particle2
			weakino = p.particle2 if is_squark(p.particle1) else p.particle1
			return 'ns',self.weakino_map[weakino],1,self.squark_map[abs(sq)],1
		if p.has_gluino() and p.has_weakino():
			return 'ng',self.weakino_map[p.particle1 if is_gluino(p.particle2) else p.particle2],1,1,1
		if p.has_gluino() and p.has_squark():
			return 'sg',1,1,self.squark_map[abs(p.particle1 if is_gluino(p.particle2) else p.particle2)],1
		if is_weakino(p.particle1) and is_weakino(p.particle2):
			return 'nn',self.weakino_map[p.particle1],self.weakino_map[p.particle2],1,1
		if is_slepton(p.particle1) and is_slepton(p.particle2):
			return 'll',self.slepton_map[(p.particle1,p.particle2)],1,1,1
		if is_gluino(p.particle1) and is_gluino(p.particle2):
			return 'gg',1,1,1,1
		if is_squark(p.particle1) and is_squark(p.particle2):
			if p.particle1 > 0 and p.particle2 >0:
				return 'ss',1,1,self.squark_map[abs(p.particle1)],self.squark_map[abs(p.particle2)]
			elif (p.particle1 > 0 and p.particle2 <0 )or(p.particle1 < 0 and p.particle2 >0 ) :
				s = p.particle1 if p.particle1 > p.particle2 else p.particle2
				b = p.particle2 if p.particle1 > p.particle2 else p.particle1
				return 'sb',1,1,self.squark_map[abs(s)],self.squark_map[abs(b)]
		return "UNKNOWN_PROCESS_OR_UNIMPLEMENTED_PROCESS"

	def _get_ps_input(self, p:Input) -> dict:
		d = {}
		d["ps_inlo"] = int(p.order)
		d["ps_final_state_in"],d["ps_ipart1_in"],d["ps_ipart2_in"],d["ps_isquark1_in"],d["ps_isquark2_in"] = self._get_ps_proc(p)
		return d
			
	def _check_input(self,p: Input)-> bool:
		"""Checks input parameter for compatibility with Prospino"""
		if p.mu_f != 1. or  p.mu_r != 1.:
			warnings.warn("Prospino2 does not support varying the scales manually.")
			return False
		if p.pdf_lo != "cteq6l1"  or  p.pdf_nlo != "cteq66":
			warnings.warn("Prospino2 does not support all pdfs (CTEQ6L1 and CTEQ66 allowed defaults).")
			return False
		if p.order != Order.LO and  p.order != Order.NLO:
			warnings.warn("Order must be one of LO/NLO in Prospino2.")
			return False
		return True

	def _is_valid(self, file: str, p: Input, d) -> bool:
		return super()._is_valid(file, p, d)

	def _parse_file(self, file: str) -> Result:
		ret=[]
		with open(file) as output:
			for line in output:
				if line.startswith("nn") or line.startswith("ng") or line.startswith("ns") or line.startswith("sg") or line.startswith("ll") or line.startswith("gg") or line.startswith("ss") or line.startswith("sb"): # TODO generalize
					for s in line[2:].split():
						ret.append(float(s))
		return ProspinoResult(ufloat(ret[8],ret[8]*ret[9]),ufloat(ret[10],ret[10]*ret[11]) if ret[10]!= 0. else None,None)	

	def _prepare(self, p: Input, **kwargs) -> RunParam:
		rp = super()._prepare(p,**kwargs)
		if not rp.skip:
			d = p.__dict__
			data = pkgutil.get_data(__name__, "prospino_main.f90_template").decode('utf-8')
			src = Template(data)
			#compute dependent pieces for template
			d["ps_inlo"] = int(p.order)
			for k,v in self._get_ps_input(p).items():
				d[k]=v
			result = src.substitute(d)
			#open(rp.in_file, "w").write(result)
			open(rp.out_file, "w").write(result + "\n\n")
			rdir = self.get_output_dir() + rp.name + ".rdir"
			if os.path.exists(rdir) and os.path.isdir(rdir):
				shutil.rmtree(rdir)
			shutil.copytree(self.get_path(),rdir)
			open(rdir  +"/prospino_main.f90", "w").write(result)
			# compile
			subprocess.Popen("cd " + rdir + " && make", shell=True,stdout=subprocess.DEVNULL).wait()

			open(rp.execute, "w").write("#!/bin/sh\n"+
				"H=$PWD && cd {rdir}&& ./prospino_2.run > tmp.out && cd $H  && cat {rdir}/tmp.out  >> {out} && cat {rdir}/prospino.dat>> {out} && rm -rf {rdir}".format(
				rdir=rdir,out=rp.out_file))
			st = os.stat(rp.execute)
			os.chmod(rp.execute, st.st_mode | stat.S_IEXEC)


			sname = d['slha']
			shutil.copy(self.get_output_dir() + sname,rdir  +"/prospino.in.les_houches")
			with open(self.get_output_dir() + sname, 'r') as f:
				open(rp.out_file, "a").write(f.read() + "\n\n")
		return rp



# Legacy 
default_prospino_runner = ProspinoRunner("~/git/Prospino2/")
"""Default Prospino Runner to provide backward compatibility"""
run = default_prospino_runner.run
set_path = default_prospino_runner.set_path
get_path = default_prospino_runner.get_path