import galsim
import numpy as np
import cmath

class FieldDistortion(object):

    def __init__(self, chip, fdModel=None, fdModel_path=None, img_rot=0.):
        if fdModel is None:
            if hasattr(chip, 'fdModel'):
                self.fdModel = chip.fdModel
            elif fdModel_path is not None:
                import pickle
                with open(fdModel_path, "rb") as f:
                    self.fdModel = pickle.load(f)
            else:
                raise ValueError("Error: no field distortion model has been specified!")
        else:
            self.fdModel = fdModel
        self.img_rot = img_rot
        self.ifdModel = self.fdModel["wave1"]
        self.ixfdModel = self.ifdModel["xImagePos"]
        self.iyfdModel = self.ifdModel["yImagePos"]
        # first-order derivatives of the global field distortion model
        self.ifx_dx    = self.ixfdModel.partial_derivative(1,0)
        self.ifx_dy    = self.ixfdModel.partial_derivative(0,1)
        self.ify_dx    = self.iyfdModel.partial_derivative(1,0)
        self.ify_dy    = self.iyfdModel.partial_derivative(0,1)
        if "residual" in self.fdModel["wave1"]:
            self.irsModel = self.fdModel["wave1"]["residual"]["ccd" + chip.getChipLabel(chipID=chip.chipID)]
            self.ixrsModel = self.irsModel["xResidual"]
            self.iyrsModel = self.irsModel["yResidual"]
            # first-order derivatives of the residual field distortion model
            self.irx_dx    = self.ixrsModel.partial_derivative(1,0)
            self.irx_dy    = self.ixrsModel.partial_derivative(0,1)
            self.iry_dx    = self.iyrsModel.partial_derivative(1,0)
            self.iry_dy    = self.iyrsModel.partial_derivative(0,1)
        else:
            self.irsModel = None

    def isContainObj_FD(self, chip, pos_img):
        xLowI, xUpI, yLowI, yUpI = self.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, img_rot=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), None
        if not img_rot:
            img_rot = np.radians(self.img_rot)
        else:
            img_rot = np.radians(img_rot)
        x, y = pos_img.x, pos_img.y
        x = self.ixfdModel(x, y)[0][0]
        y = self.iyfdModel(x, y)[0][0]
        if self.irsModel:
            # x1LowI, x1UpI, y1LowI, y1UpI = self.irsModel["interpLimit"]
            # if (x1LowI-x)*(x1UpI-x) <=0 and (y1LowI-y)*(y1UpI-y)<=0:
            #     dx = self.ixrsModel(x, y)[0][0]
            #     dy = self.iyrsModel(x, y)[0][0]
            #     x += dx
            #     y += dy
            #     # field distortion induced ellipticity
            #     ix_dx = self.ifx_dx(x, y) + self.irx_dx(x, y)
            #     ix_dy = self.ifx_dy(x, y) + self.irx_dy(x, y)
            #     iy_dx = self.ify_dx(x, y) + self.iry_dx(x, y)
            #     iy_dy = self.ify_dy(x, y) + self.iry_dy(x, y)
            # else:
            #     return galsim.PositionD(-1, -1), None
            dx = self.ixrsModel(x, y)[0][0]
            dy = self.iyrsModel(x, y)[0][0]
            x += dx
            y += dy
            # field distortion induced ellipticity
            ix_dx = self.ifx_dx(x, y) + self.irx_dx(x, y)
            ix_dy = self.ifx_dy(x, y) + self.irx_dy(x, y)
            iy_dx = self.ify_dx(x, y) + self.iry_dx(x, y)
            iy_dy = self.ify_dy(x, y) + self.iry_dy(x, y)
        else:
            ix_dx = self.ifx_dx(x, y)
            ix_dy = self.ifx_dy(x, y)
            iy_dx = self.ify_dx(x, y)
            iy_dy = self.ify_dy(x, y)
        g1k_fd   = 0.0 + (iy_dy - ix_dx) / (iy_dy + ix_dx)
        g2k_fd   = 0.0 - (iy_dx + ix_dy) / (iy_dy + ix_dx)

        # [TODO] [TESTING] Rotate the shear:
        g_abs = np.sqrt(g1k_fd**2 + g2k_fd**2)
        phi = cmath.phase(complex(g1k_fd, g2k_fd))
        # g_abs = 0.7
        g1k_fd = g_abs * np.cos(phi + 2*img_rot)
        g2k_fd = g_abs * np.sin(phi + 2*img_rot)

        fd_shear = galsim.Shear(g1=g1k_fd, g2=g2k_fd)
        return galsim.PositionD(x, y), fd_shear