import galsim
import numpy as np

class FieldDistortion(object):

    def __init__(self, fdModel=None, fdModel_path=None):
        if fdModel is None:
            import pickle
            if fdModel_path is not None:
                with open(fdModel_path, "rb") as f:
                    self.fdModel = pickle.load(f)
            else:
                # with open("/data/simudata/CSSOSDataProductsSims/data/FieldDistModelv1.0.pickle", "rb") as f:
                with open("/data/simudata/CSSOSDataProductsSims/data/FieldDistModelv2.0.pickle", "rb") as f:
                    self.fdModel = pickle.load(f)
        else:
            self.fdModel = fdModel

    def isContainObj_FD(self, chip, pos_img):
        # ifdModel = self.fdModel["ccd" + chip.getChipLabel(chipID=chip.chipID)]["wave1"]
        # x, y = pos_img.x, pos_img.y
        # xLowI, xUpI, yLowI, yUpI = ifdModel["InterpLim"]
        ifdModel = self.fdModel["wave1"]
        xLowI, xUpI, yLowI, yUpI = ifdModel["interpLimit"]
        x, y = pos_img.x, pos_img.y
        if (xLowI - x)*(xUpI - x) > 0 or (yLowI - y) * (yUpI - y) > 0:
            return False
        return True

    def get_Distorted(self, chip, pos_img, bandpass=None):
        """ Get the distored position for an undistorted image position

        Parameters:
            chip:       A 'Chip' object representing the
                        chip we want to extract PSF from.
            pos_img:    A 'galsim.Position' object representing
                        the image position.
            bandpass:   A 'galsim.Bandpass' object representing
                        the wavelength range.
        Returns:
            pos_distorted:      A 'galsim.Position' object representing
                                the distored position.
        """
        if not self.isContainObj_FD(chip=chip, pos_img=pos_img):
            return galsim.PositionD(-1, -1)
        # ifdModel = self.fdModel["ccd" + chip.getChipLabel(chipID=chip.chipID)]["wave1"]
        ifdModel = self.fdModel["wave1"]
        irsModel = self.fdModel["wave1"]["residual"]["ccd" + chip.getChipLabel(chipID=chip.chipID)]
        x, y = pos_img.x, pos_img.y
        x = ifdModel["xImagePos"](x, y)[0][0]
        y = ifdModel["yImagePos"](x, y)[0][0]
        x1LowI, x1UpI, y1LowI, y1UpI = irsModel["interpLimit"]
        if (x1LowI-x)*(x1UpI-x) <=0 and (y1LowI-y)*(y1UpI-y)<=0:
            dx = irsModel["xResidual"](x, y)[0][0]
            dy = irsModel["yResidual"](x, y)[0][0]
            x += dx
            y += dy
        else:
            return galsim.PositionD(-1, -1)
        return galsim.PositionD(x, y)