#! /usr/bin/env python3

# AutoNICER
# Copyright 2022 Nicholas Kuechel
# License Apache 2.0

import subprocess as sp
import os
import pandas as pd
import time
import numpy as np
from astroquery.heasarc import Heasarc
from astroquery import exceptions
from astropy.table import Table
from astropy.time import Time
from astropy.io import fits
from termcolor import colored
import datetime


class AutoNICER(object):
	def __init__(self, src=None):
		self.st = True
		self.xti = 0
		self.observations = []
		self.years = []
		self.months = []
		self.ras = []
		self.decs = []
		self.caldb_ver = ""
		
		print(colored("##############  Auto NICER  ##############", "cyan"))
		print()
		self.obj = src
		self.bc_sel = "y"
		self.q_set = "n"
		self.tar_sel = "y"
		self.q_path = 0
		self.q_name = 0
		if src == None:
			self.startup()
			
	def startup(self):
		"""
		Prompted setting for if autonicer is just called w/ no paramerters set
		"""
		self.obj = str(input("Target: "))
			
		self.bc_sel = str(input("Apply Bary-Center Correction: [y] "))
		if self.bc_sel == "":
			self.bc_sel = "y"
			
		self.q_set = str(input("Write Output Log: [n] "))
		self.tar_sel = str(input("Compress XTI files (.tar.gz): [y] "))
		if self.tar_sel == "":
			self.tar_sel = "y"
			
		self.q_set = self.q_set.lower()
		if self.q_set == "y":
			ne = str(input("New or Add to existing Log: "))
			if ne.lower() == "add":
				self.q_path = str(input("Input Que: "))
				if self.q_path[0] == r"'" or self.que_path[0] == r'"':
					self.q_path = self.q_path.replace("'", "")
					self.q_path = self.q_path.replace(" ", "")
					self.q_path = self.q_path.replace('"', "")
			elif ne.lower() == "new":
				self.q_name = str(input("Name of output log file (no .csv): "))
				
		else:
			self.q_set == "n"
					
	def call_nicer(self):
		"""
		Querys the nicermastr catalog for all observations of the specified source(self.obj)
		"""
		heasarc = Heasarc()
		try:
			xti = heasarc.query_object(
				self.obj, mission="nicermastr"
				)  # calls NICER master catalogue for an input object
		except exceptions.InvalidQueryError:
			print(colored(f"UNABLE TO RESOLVE {self.obj} in HEASARC!", "red"))
			print(colored(f"Exiting ...", "red"))
			exit()
		else:
			xti = Table(xti).to_pandas()
			cnt = 0
			for i in xti["OBSID"]:  # converts the form of the NICER obsid's to strings
				i = i.decode()
				xti.loc[cnt, "OBSID"] = str(i)
				cnt = cnt + 1
			cnt = 0
			for i in xti["TIME"]:  # converts times from mjd to datetime format
				t0 = Time(i, format="mjd").to_datetime()
				xti.loc[cnt, "TIME"] = t0
				cnt = cnt + 1
			self.xti = xti

	def make_cycle(self):
		"""
		Makes a Cycle# column in resulting query table from obsid table
		"""
		cycle = []
		for i in self.xti["OBSID"]:
			convo = np.floor(float(i) * 10 ** (-9))
			cycle.append(int(convo))
		self.xti["Cycle#"] = cycle
		return self.xti
		
	def get_caldb_ver(self):
		"""
		Gets most up to nicer caldb version
		"""
		caldb = sp.run("nicaldbver", shell=True, capture_output=True, encoding="utf-8")
		convo = str(caldb.stdout).split("\n")
		return convo[0]

	def sel_obs(self, enter):
		"""
		Selects and queues up the obsid desired
		"""
		selstate = True
		dup_cnt = self.observations.count(enter)
		if dup_cnt != 0:
			print(f"{enter} is already queued up... ignoring")
		else:
			try: 
				row = self.xti.loc[self.xti["OBSID"] == enter]
				print(f"Adding {enter}")
				self.observations.append(enter)
				dt = row["TIME"]
				row.reset_index(drop=True, inplace=True)
				dt.reset_index(drop=True, inplace=True)
				year = str(dt[0].year)
				self.years.append(year)
				month = dt[0].month
				# basic if else statement to fix single digit months not having a zero out front
				if month < 10:
					month = "0" + str(month)
				else:
					month = str(month)
				self.months.append(month)
				self.ras.append(row["RA"][0])
				self.decs.append(row["DEC"][0])
				return selstate
			except KeyError:
				print(colored("OBSID NOT FOUND!", "red"))
				del self.observations[-1]
				selstate = False
				return selstate

	def rm_obs(self, cmd):
		"""
		Manages the removal command types
		"""
		if cmd == "all":
			self.observations.clear()
			self.years.clear()
			self.months.clear()
			self.ras.clear()
			self.decs.clear()
		else:
			n = 0
			if cmd == "back":
				n = -1
			else:
				n = self.observations.index(cmd)
			del self.observations[n]
			del self.years[n]
			del self.months[n]
			del self.ras[n]
			del self.decs[n]
			
	def commands(self, enter):
		if enter[0].lower() == "done":
			# Command to finish selection of obs.
			self.st = False
			self.pull_reduce()
			
		elif enter[0].lower() == "sel":
			# displays all selected observations in the cmd line
			print("Observations Selected:")
			for i in self.observations:
				print(i)
			return True
			
		elif enter[0] == None or enter[0] == "":
			# Error message for nothing entered in the prompt
			print("Nothing entered...")
			print("!!!ENTER SOMETHING!!!")
			return True
			
		elif enter[0].lower() == "back":
			# Deletes the previously entered obsid
			try:
				print(f"Removing {self.observations[-1]}")
			except IndexError:
				print(colored("Nothing found to Remove!", "red"))
			else:
				self.rm_obs("back")
			return True
			
		elif enter[0].lower() == "cycle":
			row = self.make_cycle().loc[
				self.make_cycle()["Cycle#"] == float(enter[1])
				]
			for i in row["OBSID"]:
				self.sel_obs(i)
			return True
			
		elif enter[0].lower() == "rm":
			try:
				self.rm_obs(enter[1])
			except ValueError:
				print(colored("Nothing found to Remove!", "red"))
			return True
			
		elif enter[0] == "settings":
			print(f"Target: {self.obj}")
			print(f"Barycenter Correction: {self.bc_sel}")
			if self.q_set == "y":
				print(f"Log Name: {q_name}")
				print(f"Output Log: {self.q_set}")
			print(f".tar.gz compresion: {self.tar_sel}")
			
		elif enter[0] == "exit":
			exit()
			
		else:
			try:
				if int(enter[0]) > (10 ** 8):
					obsst = self.sel_obs(enter[0])
					return obsst
				else:
					raise ValueError
			except ValueError:
				print("Unknown Entry")
				return False
					
	def command_center(self, enter=None):
		"""
		Function to manage all the commands
		"""
		# prompts the user to select obs to be pulled and reduced
		orig_in = enter
		cmdstate = None
		self.st = True
		while self.st == True:
			if orig_in != None:
				cmdstate = self.commands(enter.split(" "))
				self.st = False
			else:
				cmdstate = enter = str(input(colored("autoNICER", "blue") + " > ")).split(" ")
				self.commands(enter)
			
	def nicer_compress(self):
		"""
		compresses .evt files
		"""
		print(colored("##########  .tar.gz compression  ##########", "green"))

		def tar_compr(file):
			"""
			commands for a .tar.gz compression of a single file and removal of original file after compression is complete
			"""
			if file == "":
				pass
			else:
				sp.call(f"tar czvf {file}.tar.gz {file}", shell=True)
				sp.call(f"rm -r {file}", shell=True)

		print("Compressing ufa.evt files")
		print("----------------------------------------------------------")
		# files and loop to compress the ufa files
		files = sp.run("ls *ufa.evt", shell=True, capture_output=True, encoding="utf-8")
		for i in str(files.stdout).split("\n"):
			tar_compr(i)

		# compression of the non-bc mpu7_cl.evt file if barycenter correction is selected
		if self.bc_sel.lower() == "n":
			pass
		else:
			print("")
			print("Compressing cl.evt files")
			print("----------------------------------------------------------")
			# files and liip to compress the cl files
			cl_file = sp.run(
				"ls ni*cl.evt", shell=True, capture_output=True, encoding="utf-8"
			)
			for i in str(cl_file.stdout).split("\n"):
				tar_compr(i)
				
	def add2q(self, q, base_dir, obsid):
		"""
		Adds processed data to existing log
		"""
		newline = pd.DataFrame({
				"Input":[f"{base_dir}/{obsid}/xti/event_cl/bc{obsid}_0mpu7_cl.evt"],
				"OBSID":[f"NI{obsid}"],
				"CALDB_ver": [f"{self.caldb_ver}"],
				"DateTime": [f"{datetime.datetime.now()}"]
				},
		)
		q = pd.concat([q, newline])
		q.to_csv(self.q_path, index=False)
		
	def reduce(self, obsid):
		"""
		Performs standardized data reduction scheme calling nicerl2 and barycorr if set
		"""
		sp.call(f"nicerl2 indir={obsid}/ clobber=yes", shell=True)
		obsindex = self.observations.index(obsid)
		if self.bc_sel.lower() == "n":
			pass
		else:
			sp.call(
					f"barycorr infile={obsid}/xti/event_cl/ni{obsid}_0mpu7_cl.evt outfile={obsid}/xti/event_cl/bc{obsid}_0mpu7_cl.evt 						orbitfiles={obsid}/auxil/ni{obsid}.orb refframe=ICRS ra={self.ras[obsindex]} dec={self.decs[obsindex]} ephem=JPLEPH.430",
					shell=True,
				)

	def pull_reduce(self):
		"""
		Downloads the NICER data
		Puts the retrieved data through a standardized data reduction scheme
		"""
		downCommand = (
			"wget -q -nH --no-check-certificate --cut-dirs=5 -r -l0 -c -N -np -R "
			+ "'index*'"
			+ " -erobots=off --retr-symlinks https://heasarc.gsfc.nasa.gov/FTP/nicer/data/obs/"
		)

		for obsid in self.observations:
			qindex = self.observations.index(obsid)
			print("")
			print("--------------------------------------------------------------")
			print("             Prosessing OBSID: " + colored(str(obsid), "cyan"))
			print("--------------------------------------------------------------")
			pull_templ = (
				f"{downCommand}{self.years[qindex]}_{self.months[qindex]}//{obsid}"
			)
			end_args = f"--show-progress --progress=bar:force"
			print(colored("Downloading xti data...", "green"))
			sp.call(f"{pull_templ}/xti/ {end_args}", shell=True)
			print(colored("Downloadng log data...", "green"))
			sp.call(f"{pull_templ}/log/ {end_args}", shell=True)
			print(colored("Downloading auxil data...", "green"))
			sp.call(f"{pull_templ}/auxil/ {end_args}", shell=True)
			self.caldb_ver = self.get_caldb_ver()
			self.reduce(obsid)

			base_dir = os.getcwd()
			os.chdir(f"{obsid}/xti/event_cl/")
			if self.tar_sel == "n" or self.tar_sel == "N":
				pass
			else:
				self.nicer_compress()
			os.chdir(base_dir)
			if self.q_set == "y" and self.q_path != 0:
				read_q = pd.read_csv(self.q_path)
				self.add2q(read_q, base_dir, obsid)
			elif self.q_set == "y" and self.q_path == 0:
				q = pd.DataFrame({"Input":[], "Name":[], "CALDB_ver": [], "DateTime": []})
				self.q_path = f"{base_dir}/{self.q_name}.csv"
				self.add2q(q, base_dir, obsid)
			else:
				pass
		
