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, 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
import pyslha
from uncertainties import ufloat



class NLLfastRunner(Runner):
	#TODO treat stop sbot separately
	#TODO separate nll and nnll

	def _get_nf_proc(self,p:Input):
		d = pyslha.read(self.get_output_dir() + p.slha)
		mg = d.blocks["MASS"][1000021]
		ms = 0.
		for r in range(1000001,1000007):
			ms += d.blocks["MASS"][r]
		for r in range(2000001,2000007):
			ms += d.blocks["MASS"][r]
		ms /= 12
		if p.has_gluino() and p.has_squark():
			return 'sg','cteq',ms,mg
		if is_gluino(p.particle1) and is_gluino(p.particle2):
			return 'gg','cteq',ms,mg
		if is_squark(p.particle1) and is_squark(p.particle2):
			if p.particle1 > 0 and p.particle2 >0:
				return 'ss','cteq',ms,mg
			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','cteq',ms,mg
		return "UNKNOWN_PROCESS_OR_UNIMPLEMENTED_PROCESS"

	def _get_nf_input(self, p:Input) -> dict:
		# TODO return masses of squark and gluino
		d = {}
		#d["ps_inlo"] = int(p.order)
		d["nf_final_state_in"],d["nf_pdf"],d["nf_squark_mass"],d["nf_gluino_mass"] = 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("NLL-fast does not support varying the scales manually.")
			return False
		if p.pdf_lo != "cteq6l1"  or  p.pdf_nlo != "cteq66":
			warnings.warn("NLL-fast does not support all pdfs (CTEQ6L1 and CTEQ66 allowed defaults).")
			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:
		#TODO parse result
		pass
		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 Result(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
			for k,v in self._get_nf_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"+
				"{path}/nll-fast {proc} {pdf} {sq} {gl}> {out}".format(
				path=self.get_path(),out=rp.out_file , proc = d["nf_final_state_in"], sq = d["nf_squark_mass"], gl = d["nf_gluino_mass"], pdf = d["nf_pdf"]))
			st = os.stat(rp.execute)
			os.chmod(rp.execute, st.st_mode | stat.S_IEXEC)


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



# Legacy 
default_nllfast_runner = NLLfastRunner("~/git/nll-fast/")
"""Default Prospino Runner to provide backward compatibility"""
run = default_nllfast_runner.run
set_path = default_nllfast_runner.set_path
get_path = default_nllfast_runner.get_path