from .FocalPlane import FocalPlane
import galsim
import os
import numpy as np
from astropy.table import Table

class Chip(FocalPlane):
	def __init__(self, chipID, ccdEffCurve_dir, config=None, treering_func=None):
		# Get focal plane (instance of paraent class) info
		# TODO: use chipID to config individual chip?
		super().__init__()
		if config is not None:
			self.npix_x 	= config["npix_x"]
			self.npix_y 	= config["npix_y"]
			self.read_noise = config["read_noise"]
			self.dark_noise = config["dark_noise"]
			self.pix_scale 	= config["pix_scale"]
			self.gain		= config["gain"]
		else:
			# Default setting
			self.npix_x 	= 9232
			self.npix_y 	= 9216
			self.read_noise = 5.0	# e/pix
			self.dark_noise = 0.02	# e/pix/s
			self.pix_scale 	= 0.074	# pixel scale
			self.gain		= 1.0

		# A chip ID must be assigned
		self.chipID = int(chipID)

		# Get corresponding filter info
		self.filter_id, self.filter_type = self.getChipFilter()

		# Get boundary (in pix)
		self.bound = self.getChipLim()
		self.ccdEffCurve_dir = ccdEffCurve_dir
		self.effCurve = self._getChipEffCurve(self.filter_type)

		# Define the sensor
		# self.sensor = galsim.Sensor()
		self.sensor = galsim.SiliconSensor(strength=10., treering_func=treering_func)

	def _getChipEffCurve(self, filter_type):
		if filter_type in ['nuv', 'u']: filename = 'UV0.txt'
		if filter_type in ['g', 'r']: filename = 'Astro_MB.txt'
		if filter_type in ['i', 'z', 'y']: filename = 'Basic_NIR.txt'
		path = os.path.join(self.ccdEffCurve_dir, filename)
		table = Table.read(path, format='ascii')
		throughput = galsim.LookupTable(x=table['col1'], f=table['col2'], interpolant='linear')
		bandpass = galsim.Bandpass(throughput, wave_type='nm')
		return bandpass

	def getChipFilter(self, chipID=None, filter_layout=None):
		"""Return the filter index and type for a given chip #(chipID)
		"""
		filter_type_list = ["nuv","u", "g", "r", "i","z","y"]
		# TODO: maybe a more elegent way other than hard coded?
		# e.g. use something like a nested dict:
		if filter_layout is not None:
			return filter_layout[chipID][0], filter_layout[chipID][1]
		if chipID == None:
			chipID = self.chipID

		# updated configurations
		if chipID>30 or chipID<1: raise ValueError("!!! Chip ID: [1,30]")
		if chipID in [10, 11, 20, 21]: filter_type = 'y'
		if chipID in [15, 16]:         filter_type = "z"
		if chipID in [9, 22]:			filter_type = "i"
		if chipID in [12, 19]:         filter_type = "u"
		if chipID in [7, 24]:         filter_type = "r"
		if chipID in [14, 13, 18, 17]:    filter_type = "nuv"
		if chipID in [8, 23]:         filter_type = "g"
		if chipID in [6, 5, 25, 26]:	filter_type = "GI"
		if chipID in [27, 30, 1, 4]:	filter_type = "GV"
		if chipID in [28, 29, 2, 3]:	filter_type = "GU"
		filter_id = filter_type_list.index(filter_type)

		return filter_id, filter_type

	def getChipLim(self, chipID=None):
		"""Calculate the edges in pixel for a given CCD chip on the focal plane
		NOTE: There are 5*4 CCD chips in the focus plane for photometric observation.
		Parameters:
			chipID:			int
							the index of the chip
		Returns:
			A galsim BoundsD object
		"""
		if chipID == None:
			chipID = self.chipID
		
		gx = self.npix_gap_x
		gy1, gy2 = self.npix_gap_y

		# xlim of a given ccd chip
		xrem = (chipID-1)%self.nchip_x - self.nchip_x // 2
		xcen = (self.npix_x + gx) * xrem
		nx0 = xcen - self.npix_x//2 + 1
		nx1 = xcen + self.npix_x//2

		# ylim of a given ccd chip
		yrem = 2*((chipID-1)//self.nchip_x) - (self.nchip_y-1)
		ycen = (self.npix_y//2 + gy1//2) * yrem
		if chipID <= 6: ycen = (self.npix_y//2 + gy1//2) * yrem - (gy2-gy1)
		if chipID >= 25: ycen = (self.npix_y//2 + gy1//2) * yrem + (gy2-gy1)
		ny0 = ycen - self.npix_y//2 + 1
		ny1 = ycen + self.npix_y//2

		return galsim.BoundsD(nx0-1, nx1-1, ny0-1, ny1-1)


	def getSkyCoverage(self, wcs):
		return super().getSkyCoverage(wcs, self.bound.xmin, self.bound.xmax, self.bound.ymin, self.bound.ymax)


	def getSkyCoverageEnlarged(self, wcs, margin=0.5):
		"""The enlarged sky coverage of the chip
		"""
		margin /= 60.0
		bound = self.getSkyCoverage(wcs)
		return galsim.BoundsD(bound.xmin - margin, bound.xmax + margin, bound.ymin - margin, bound.ymax + margin)


	def isContainObj(self, ra_obj, dec_obj, sky_coverage):
		if (ra_obj - sky_coverage.xmin) * (ra_obj - sky_coverage.xmax) > 0.0 or (dec_obj - sky_coverage.ymin) * (dec_obj - sky_coverage.ymax) > 0.0:
			return False
		return True


	def getChipNoise(self, exptime=150.0):
		# TODO: put it here or make it a separate class?
		noise = self.dark_noise * exptime + self.read_noise**2
		return noise


	def addNoise(self, img, exptime=150.0, sky_noise=0., seed=31415):
		rng = galsim.BaseDeviate(seed)
		dark_noise = galsim.DeviateNoise(galsim.PoissonDeviate(rng, self.dark_noise*exptime))
		img.addNoise(dark_noise)
		ccd_noise = galsim.CCDNoise(rng, sky_level=sky_noise, gain=self.gain, read_noise=self.read_noise)
		img.addNoise(ccd_noise)
		return img