# Extract flyby parameters from CASSINI

# Built in Python functions
from datetime import datetime, timedelta
import os
import zipfile

# External Python libraries (installed via pip install)
from bs4 import BeautifulSoup
import logging
import pandas as pd
from urllib import request, error

# Internal Pydar reference to access functions, global variables, and error handling
import pydar

## Logging set up for .INFO
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
stream_handler = logging.StreamHandler()
logger.addHandler(stream_handler)

resolution_types = ["B", "D", "F", "H", "I"] # 2, 8, 32, 128, 256 pixels/degree
datafile_types_columns = ["ABDR", "ASUM", "BIDR", "LBDR", "SBDR", "STDR"]

def getFlybyData():
	# Header: Titan flyby id, Radar Data Take Number, Sequence number, Orbit Number/ID
	flyby_id = []
	flyby_radar_take_num = []
	flyby_csv_file = os.path.join(os.path.dirname(__file__), 'data', 'cassini_flyby.csv')  # get file's directory, up one level, /data/*.csv
	flyby_dataframe = pd.read_csv(flyby_csv_file)
	for index, row in flyby_dataframe.iterrows():
		row = row.tolist()
		flyby_id.append(row[0])
		radar_take_num = row[1].split(" ")[1]
		while len(radar_take_num) < 4:
			radar_take_num = "0" + radar_take_num # set all radar take numbers to be four digits long: 229 -> 0229
		flyby_radar_take_num.append(radar_take_num)
	# returns a list of flyby IDs and associated Radar Data Take Number
	return flyby_id, flyby_radar_take_num

def convertFlybyIDToObservationNumber(flyby_id=None):
	# convert Flyby ID to Observation Number to find data files
	pydar.errorHandlingConvertFlybyIDToObservationNumber(flyby_id=flyby_id)

	flyby_csv_file = os.path.join(os.path.dirname(__file__), 'data', 'cassini_flyby.csv')  # get file's directory, up one level, /data/*.csv
	flyby_dataframe = pd.read_csv(flyby_csv_file)
	for index, row in flyby_dataframe.iterrows():
		if row.iloc[0] == flyby_id:
			observation_number = row.iloc[1].split(" ")[1]
			while len(observation_number) < 4:
				observation_number = "0" + observation_number # set all radar take numbers to be four digits long: 229 -> 0229
			return observation_number

def convertObservationNumberToFlybyID(flyby_observation_num=None):
	# convert Flyby ID to Observation Number to find data files
	if flyby_observation_num is not None:
		if type(flyby_observation_num) != str:
			raise ValueError(f"[flyby_observation_num]: Must be a str, current type = '{type(flyby_observation_num)}'")
		else:
			while len(flyby_observation_num) < 4:
				flyby_observation_num = "0" + flyby_observation_num # set all radar take numbers to be four digits long: 229 -> 0229

	pydar.errorHandlingConvertObservationNumberToFlybyID(flyby_observation_num=flyby_observation_num)

	flyby_csv_file = os.path.join(os.path.dirname(__file__), 'data', 'cassini_flyby.csv')  # get file's directory, up one level, /data/*.csv
	flyby_dataframe = pd.read_csv(flyby_csv_file)
	for index, row in flyby_dataframe.iterrows():
		take_ob_num = "0" + row.iloc[1].split(" ")[1]
		if take_ob_num == flyby_observation_num:
			return row.iloc[0] # returns flyby ID

def retrieveJPLCoradrOptions():
	# Read JPL Options from CSV
	coradr_csv_file = os.path.join(os.path.dirname(__file__), 'data', 'coradr_jpl_options.csv')  # get file's directory, up one level, /data/*.csv
	coradr_dataframe = pd.read_csv(coradr_csv_file)
	return coradr_dataframe

def retrieveMostRecentVersionNumber(flyby_observiation_num=None):
	# Return the CORADAR value with the most recent version from a list of possible options
	coradr_dataframe = retrieveJPLCoradrOptions()
	jpl_coradr_options = []
	for index, row in coradr_dataframe.iterrows():
		row = row.tolist()
		jpl_coradr_options.append(row[0])

	find_cordar_listing = f"CORADR_{flyby_observiation_num}"
	version_types_available = list(filter(lambda x: find_cordar_listing in x, jpl_coradr_options))
	more_accurate_model_number = version_types_available[-1] # always choose the last and more up to date version number (Currently, v3)
	logger.info(f"Most recent CORADR version is {more_accurate_model_number} from the available list {version_types_available}")
	return more_accurate_model_number

def retrieveCoradrWithoutBIDR():
	# Return a list of valid flyby observation numbers that do not contain BIDR
	coradr_dataframe = retrieveJPLCoradrOptions()
	coradar_without_bidr = []
	for index, row in coradr_dataframe.iterrows():
		if row["Is a Titan Flyby"]: # check only flybys that are valid Titan flybys
			if "V" not in row["CORADR ID"] and row["Contains BIDR"] is False: # if not a version row, and does not contain BIDR
				coradar_without_bidr.append(row["CORADR ID"].split("_")[1]) # store the observation number from the CORADR
	return coradar_without_bidr

def downloadAAREADME(cordar_file_name, segment_id):
	# Download AAREADME.txt within a CORADR directory
	aareadme_name = "AAREADME.TXT"
	aareadme_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/{aareadme_name}"

	# Retrieve a list of all elements from the base URL to download AAREADME.txt
	logger.info(f"Retrieving {cordar_file_name} {aareadme_name}")
	aareadme_name = os.path.join(f"pydar_results/{cordar_file_name}_{segment_id}", aareadme_name)
	try:
		request.urlretrieve(aareadme_url)
	except error.HTTPError as err:
		raise HTTPError(f"Unable to access: {aareadme_url}\nError (and exiting): '{err.code}'")
	else:
		response = request.urlretrieve(aareadme_url, aareadme_name)

def downloadBIDRCORADRData(cordar_file_name, segment_id, resolution_px):
	# Download BDIR files
	base_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/BIDR/"
	logger.info(f"Retrieving BIDR filenames from: {base_url}\n")

	# Retrieve a list of all elements from the base URL to download
	base_html = request.urlopen(base_url).read()
	soup = BeautifulSoup(base_html, 'html.parser')
	table = soup.find('table', {"id": "indexlist"})
	table_text = (table.text).split("\n")
	url_filenames = []
	all_bidr_files = []
	for txt in table_text:
		if txt.startswith('BI'):
			filename = (txt.split('/')[0]).split(".")[0]
			if 'LBL' in (txt.split('/')[0]).split(".")[1]:
				filename += '.LBL'
			if 'ZIP' in (txt.split('/')[0]).split(".")[1]:
				filename += '.ZIP'
			all_bidr_files.append(filename)
			if segment_id in filename: # only save certain segments
				for resolution in resolution_px: # only save top x resolutions
					bi_types = ["B", "E", "T", "N", "M", "L"] # BI<OPTION>Q<RESOLUTION>
					for bi in bi_types:
						if f"BI{bi}Q{resolution}" in filename:
							url_filenames.append(filename)

	logger.info(f"All BIDR files found with specified resolution, segment, and flyby identification: {url_filenames}\n")
	if len(url_filenames) == 0:
		raise ValueError(f"No BIDR files found with resolution, segment, and flyby identification. Please use different parameters to retrieve data.\nAll files found: {all_bidr_files}")

	for i, coradr_file in enumerate(url_filenames):
		if 'LBL' in coradr_file:
			label_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/BIDR/{coradr_file}"
			logger.info(f"Retrieving [{i+1}/{len(url_filenames)}]: {label_url}")
			label_name = label_url.split("/")[-1].split(".")[0] + ".LBL"
			label_name = os.path.join(f"pydar_results/{cordar_file_name}_{segment_id}", label_name)
			try:
				request.urlretrieve(label_url)
			except error.HTTPError as err:
				raise HTTPError(f"Unable to access: {label_url}\nError (and exiting): '{err.code}'")
			else:
				response = request.urlretrieve(label_url, label_name)
		if 'ZIP' in coradr_file:
			data_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/BIDR/{coradr_file}"
			logger.info(f"Retrieving [{i+1}/{len(url_filenames)}]: {data_url}")
			zipfile_name = data_url.split("/")[-1].split(".")[0] + ".zip"
			zipfile_name = os.path.join(f"pydar_results/{cordar_file_name}_{segment_id}", zipfile_name)
			try:
				request.urlretrieve(data_url)
			except error.HTTPError as err:
				raise HTTPError(f"Unable to access: {data_url}\nError (and exiting): '{err.code}'")
			else:
				response = request.urlretrieve(data_url, zipfile_name)
				zipped_image = zipfile_name.split(".")[0] + ".IMG"
				with zipfile.ZipFile(zipfile_name, 'r') as zip_ref:
					zipped_image_path = os.path.join(f"pydar_results/{cordar_file_name}_{segment_id}")
					zip_ref.extractall(zipped_image_path)

def downloadSBDRCORADRData(cordar_file_name, segment_id):
	# Download SBDR files
	base_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/SBDR/"
	logger.info(f"\nRetrieving SBDR filenames from: {base_url}")

	# Retrieve a SBDR file from filename at SBDR URL
	base_html = request.urlopen(base_url).read()
	soup = BeautifulSoup(base_html, 'html.parser')
	table = soup.find('table', {"id": "indexlist"})
	table_text = (table.text).split("\n")
	sbdr_files = []
	sbdr_filename = ''
	for txt in table_text:
		if txt.startswith('SBDR'):
			sbdr_filename = (txt.split('/')[0]).split(".")[0]
			if 'TAB' in (txt.split('/')[0]).split(".")[1]: # TAB information
				sbdr_filename += '.TAB'
			if 'FMT' in (txt.split('/')[0]).split(".")[1]: # FMT information required to read TAB
				sbdr_filename += '.FMT'
			sbdr_files.append(sbdr_filename)

	logger.info("SBDR files found: {0}".format(sbdr_files))
	if len(sbdr_files) == 0:
		raise ValueError("No SBDR files were found with resolution, segment, and flyby identification. Please use different parameters to retrieve data")

	for sbdr_file in sbdr_files:
		sbdr_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/SBDR/{sbdr_file}"
		logger.info(f"Retrieving SBDR file '{sbdr_file}': {sbdr_url}")
		sbdr_name = os.path.join(f"pydar_results/{cordar_file_name}_{segment_id}", sbdr_file)
		try:
			request.urlretrieve(sbdr_url)
		except error.HTTPError as err:
			raise HTTPError(f"Unable to access: {sbdr_url}\nError (and exiting): '{err.code}'")
		else:
			response = request.urlretrieve(sbdr_url, sbdr_name)

def downloadAdditionalDataTypes(cordar_file_name, segment_id, additional_data_type):
	# Download additional data types
	additional_data_url = f"https://planetarydata.jpl.nasa.gov/img/data/cassini/cassini_orbiter/{cordar_file_name}/DATA/{additional_data_type}"
	logger.info(f"\n[TODO: does not currently download] '{additional_data_type}': {additional_data_url}")
	# TODO: add functionality for which files should be downloaded
	# This function does not currently have functionality in pydar

def extractFlybyDataImages(flyby_observation_num=None,
							flyby_id=None,
							segment_num=None,
							additional_data_types_to_download=[],
							resolution='I',
							top_x_resolutions=None):

	if flyby_id is not None and type(flyby_id) == str:
		flyby_id = flyby_id.upper() # ensure that observation number set to capitalized 'T'
	if flyby_observation_num is not None and type(flyby_observation_num) == str:
		while len(flyby_observation_num) < 4:
			flyby_observation_num = "0" + flyby_observation_num # set all radar take numbers to be four digits long: 229 -> 0229
	if top_x_resolutions is not None:
		#logger.info("\nINFO: [top_x_resolutions] in use, overriding resolution '{0}' to save the top {1} resolutions".format(resolution, top_x_resolutions))
		resolution = None # set default resolution to None if selecting the top x resolutions

	# Error handling:
	pydar.errorHandlingExtractFlybyDataImages(flyby_observation_num=flyby_observation_num,
											flyby_id=flyby_id,
											segment_num=segment_num,
											additional_data_types_to_download=additional_data_types_to_download,
											resolution=resolution,
											top_x_resolutions=top_x_resolutions)

	logger.debug(f"flyby_observation_num = {flyby_observation_num}")
	logger.debug(f"flyby_id = {flyby_id}")
	logger.debug(f"segment_num = {segment_num}")
	logger.debug(f"additional_data_types_to_download = {additional_data_types_to_download}")
	logger.debug(f"resolution = {resolution}")
	logger.debug(f"top_x_resolutions = {top_x_resolutions}")

	download_files = True # for debugging, does not always download files before running data

	if flyby_id is not None:  # convert flyby Id to an Observation Number
		flyby_observation_num = convertFlybyIDToObservationNumber(flyby_id)

	# Data gaps and problems from the original downlinking and satellite location, report some special cases to user
	no_associated_bidr_values = retrieveCoradrWithoutBIDR() # currently: ["0048", "0186", "0189", "0209", "0234"]
	if flyby_observation_num in no_associated_bidr_values:
		logger.info("\nINFO: due to data gaps or issues with downlinking, flyby does not have not have associated BIDR data.")
		if flyby_observation_num == "0048":
			logger.info("0048 (T4) did not have SAR data, only scatterometry and radiometry\n")
		elif flyby_observation_num == "0186":
			logger.info("0186 (T52) only has radiometery and compressed scatterometry\n")
		elif flyby_observation_num == "0189":
			logger.info("0189 (T53) only has radiometery and compressed scatterometry\n")
		elif flyby_observation_num == "0209":
			logger.info("0209 (T63) only has scatterometry and radiometry\n")
		elif flyby_observation_num == "0234":
			logger.info("0234 (T80) only has scatterometry and radiometry\n")
		else:
			logger.info(f"{flyby_observation_num} does not BIDR data\n") # possible catch for new files found without BIDR

	available_flyby_id, available_observation_numbers = getFlybyData()

	if flyby_observation_num not in available_observation_numbers:
		raise ValueError(f"Observation number '{flyby_observation_num}' NOT FOUND in available observation numbers: {available_observation_numbers}\n")
	else:
		logger.debug(f"Observation number '{flyby_observation_num}' FOUND in available observation numbers: {available_observation_numbers}\n")

	# Download information from pds-imaging site for CORADR
	flyby_observation_cordar_name = retrieveMostRecentVersionNumber(flyby_observation_num)
	if not os.path.exists('pydar_results'): os.makedirs('pydar_results')
	if not os.path.exists(f"pydar_results/{flyby_observation_cordar_name}_{segment_num}"): 
		os.makedirs(f"pydar_results/{flyby_observation_cordar_name}_{segment_num}")

	if download_files:
		# Download AAREADME.TXT
		downloadAAREADME(flyby_observation_cordar_name, segment_num)
		
		# Download BIDR
		if flyby_observation_num not in no_associated_bidr_values: # only attempt to download BIDR files for flybys that have BIDR files
			if top_x_resolutions is not None:
				downloadBIDRCORADRData(flyby_observation_cordar_name, segment_num, resolution_types[-top_x_resolutions:])
			else:
				downloadBIDRCORADRData(flyby_observation_cordar_name, segment_num, resolution)

		# Download SBDR
		downloadSBDRCORADRData(flyby_observation_cordar_name, segment_num)

		# Download additional data types (TODO)
		for data_type in additional_data_types_to_download:
			if data_type not in ["BIDR", "SBDR"]: # ignore data files that have already been downloaded
				downloadAdditionalDataTypes(flyby_observation_cordar_name, segment_num, data_type)

	# No valid parameters given, empty file
	if len(os.listdir(f"pydar_results/{flyby_observation_cordar_name}_{segment_num}")) == 0:
		logger.critical(f"\npydar_results/{flyby_observation_cordar_name}_{segment_num} is empty. Unable to find any data files with current parameters")
