#from .. mappings import DEGTORAD, RADTODEG
from hstwcs.mappings import DEGTORAD, RADTODEG
import numpy
from numpy import math 
from math import sin, sqrt, pow, cos, asin, atan2,pi
from hstwcs import utils

class MakeWCS(object):
    """
    Purpose
    =======
    Recompute WCS keywords based on PA_V3 and distortion model.
    
    Algorithm
    =========
    - update reference chip wcs
    - update extension wcs
    - update extension header
    
    """
 
    def updateWCS(cls, ext_wcs, ref_wcs):    
        """
        recomputes the basic WCS kw 
        """
        ltvoff, offshift = cls.getOffsets(ext_wcs)
            
        cls.uprefwcs(ext_wcs, ref_wcs, ltvoff, offshift)
        cls.upextwcs(ext_wcs, ref_wcs, ltvoff, offshift)
        
        kw2update = {'CD1_1': ext_wcs.wcs.cd[0,0],
                    'CD1_2': ext_wcs.wcs.cd[0,1],
                    'CD2_1': ext_wcs.wcs.cd[1,0],
                    'CD2_2': ext_wcs.wcs.cd[1,1],
                    'CRVAL1': ext_wcs.wcs.crval[0],
                    'CRVAL2': ext_wcs.wcs.crval[1],
                    'CRPIX1': ext_wcs.wcs.crpix[0],
                    'CRPIX2': ext_wcs.wcs.crpix[1],
            }
        return kw2update
    
    updateWCS = classmethod(updateWCS)
    
    def upextwcs(self, ext_wcs, ref_wcs, loff, offsh):
        """
        updates an extension wcs
        """  
        ltvoffx, ltvoffy = loff[0], loff[1]
        offshiftx, offshifty = offsh[0], offsh[1]
        ltv1 = ext_wcs.ltv1
        ltv2 = ext_wcs.ltv2
        
        if ltv1 != 0. or ltv2 != 0.:
            offsetx = backup_wcs['CRPIX1'] - ltv1 - ext_wcs.idcmodel.refpix['XREF']
            offsety = backup_wcs['CRPIX2'] - ltv2 - ext_wcs.idcmodel.refpix['YREF']
            fx,fy = ext_wcs.idcmodel.shift(ext_wcs.idcmodel.cx,ext_wcs.idcmodel.cy,offsetx,offsety)
        else:
            fx, fy = ext_wcs.idcmodel.cx, ext_wcs.idcmodel.cy
        
        ridcmodel = ref_wcs.idcmodel
        v2 = ext_wcs.idcmodel.refpix['V2REF']
        v3 = ext_wcs.idcmodel.refpix['V3REF']
        v2ref = ref_wcs.idcmodel.refpix['V2REF']

        v3ref = ref_wcs.idcmodel.refpix['V3REF']
        R_scale = ref_wcs.idcmodel.refpix['PSCALE']/3600.0
        off = sqrt((v2-v2ref)**2 + (v3-v3ref)**2)/(R_scale*3600.0)
        if v3 == v3ref:
           theta=0.0
        else:
           theta = atan2(ext_wcs.parity[0][0]*(v2-v2ref), ext_wcs.parity[1][1]*(v3-v3ref))
        
        if ref_wcs.idcmodel.refpix['THETA']: theta += ref_wcs.idcmodel.refpix['THETA']*pi/180.0

        dX=(off*sin(theta)) + offshiftx
        dY=(off*cos(theta)) + offshifty
        px = numpy.array([[dX,dY]])
        newcrval = ref_wcs.all_pix2sky_fits(px)[0]
        newcrpix = numpy.array([ext_wcs.idcmodel.refpix['XREF'] + ltvoffx, 
                                ext_wcs.idcmodel.refpix['YREF'] + ltvoffy])
        ext_wcs.wcs.crval = newcrval
        ext_wcs.wcs.crpix = newcrpix
        
        # Account for subarray offset
        # Angle of chip relative to chip
        if ext_wcs.idcmodel.refpix['THETA']:
           dtheta = ext_wcs.idcmodel.refpix['THETA'] - ref_wcs.idcmodel.refpix['THETA']
        else:
           dtheta = 0.0
        
        
        # Create a small vector, in reference image pixel scale
        # There is no parity effect here ???
        
        delXX=fx[1,1]/R_scale/3600.
        delYX=fy[1,1]/R_scale/3600.
        delXY=fx[1,0]/R_scale/3600.
        delYY=fy[1,0]/R_scale/3600.

        # Convert to radians
        rr=dtheta*pi/180.0
        
        # Rotate the vectors
        dXX = cos(rr)*delXX - sin(rr)*delYX
        dYX = sin(rr)*delXX + cos(rr)*delYX

        dXY = cos(rr)*delXY - sin(rr)*delYY
        dYY = sin(rr)*delXY + cos(rr)*delYY
        px = numpy.array([[dX+dXX,dY+dYX]])
        
        # Transform to sky coordinates
        wc = ref_wcs.all_pix2sky_fits(px)
        
        a = wc[0,0]
        b = wc[0,1]
        px = numpy.array([[dX+dXY,dY+dYY]])
        
        wc = ref_wcs.all_pix2sky_fits(px)
        c = wc[0,0]
        d = wc[0,1]

        # Calculate the new CDs and convert to degrees
        cd11 = utils.diff_angles(a,newcrval[0])*cos(newcrval[1]*pi/180.0)
        cd12 = utils.diff_angles(c,newcrval[0])*cos(newcrval[1]*pi/180.0)
        cd21 = utils.diff_angles(b,newcrval[1])
        cd22 = utils.diff_angles(d,newcrval[1])
        cd = numpy.array([[cd11, cd12], [cd21, cd22]])
        ext_wcs.wcs.cd = cd  #???
        
    upextwcs = classmethod(upextwcs)
        
    def uprefwcs(self, ext_wcs, ref_wcs, loff, offsh):
        """
        Updates the reference chip
        """
        ltvoffx, ltvoffy = loff[0], loff[1]
        offshift = offsh
        dec = ref_wcs.wcs.crval[1]
        # Get an approximate reference position on the sky
        rref = numpy.array([[ref_wcs.idcmodel.refpix['XREF']+ltvoffx, 
                            ref_wcs.idcmodel.refpix['YREF']+ltvoffy]])
        
        crval = ref_wcs.all_pix2sky_fits(rref)
        # Convert the PA_V3 orientation to the orientation at the aperture
        # This is for the reference chip only - we use this for the
        # reference tangent plane definition
        # It has the same orientation as the reference chip
        pv = troll(ext_wcs.pav3,dec,ref_wcs.idcmodel.refpix['V2REF'],
                    ref_wcs.idcmodel.refpix['V3REF'])
        # Add the chip rotation angle
        if ref_wcs.idcmodel.refpix['THETA']:
            pv += ref_wcs.idcmodel.refpix['THETA']
            
        
        # Set values for the rest of the reference WCS
        ref_wcs.wcs.crval = crval[0]
        ref_wcs.wcs.crpix = numpy.array([0.0,0.0])+offsh
        parity = ref_wcs.parity
        R_scale = ref_wcs.idcmodel.refpix['PSCALE']/3600.0
        cd11 = parity[0][0] *  cos(pv*pi/180.0)*R_scale
        cd12 = parity[0][0] * -sin(pv*pi/180.0)*R_scale
        cd21 = parity[1][1] *  sin(pv*pi/180.0)*R_scale
        cd22 = parity[1][1] *  cos(pv*pi/180.0)*R_scale
        
        rcd = numpy.array([[cd11, cd12], [cd21, cd22]])
        ref_wcs.wcs.cd = rcd
        ref_wcs.wcs.set()
        
    uprefwcs = classmethod(uprefwcs)
        

    def getOffsets(cls, ext_wcs):
        ltv1 = ext_wcs.ltv1
        ltv2 = ext_wcs.ltv2
        
        offsetx = ext_wcs.wcs.crpix[0] - ltv1 - ext_wcs.idcmodel.refpix['XREF']
        offsety = ext_wcs.wcs.crpix[1] - ltv2 - ext_wcs.idcmodel.refpix['YREF']
        
        shiftx = ext_wcs.idcmodel.refpix['XREF'] + ltv1
        shifty = ext_wcs.idcmodel.refpix['YREF'] + ltv2
        if ltv1 != 0. or ltv2 != 0.:
            ltvoffx = ltv1 + offsetx
            ltvoffy = ltv2 + offsety
            offshiftx = offsetx + shiftx
            offshifty = offsety + shifty
        else:
            ltvoffx = 0.
            ltvoffy = 0.
            offshiftx = 0.
            offshifty = 0.
            
        ltvoff = numpy.array([ltvoffx, ltvoffy])
        offshift = numpy.array([offshiftx, offshifty])
        return ltvoff, offshift
    
    getOffsets = classmethod(getOffsets)
    

def troll(roll, dec, v2, v3):
    """ Computes the roll angle at the target position based on:
            the roll angle at the V1 axis(roll),
            the dec of the target(dec), and
            the V2/V3 position of the aperture (v2,v3) in arcseconds.

        Based on the algorithm provided by Colin Cox that is used in
        Generic Conversion at STScI.
    """
    # Convert all angles to radians
    _roll = DEGTORAD(roll)
    _dec = DEGTORAD(dec)
    _v2 = DEGTORAD(v2 / 3600.)
    _v3 = DEGTORAD(v3 / 3600.)

    # compute components
    sin_rho = sqrt((pow(sin(_v2),2)+pow(sin(_v3),2)) - (pow(sin(_v2),2)*pow(sin(_v3),2)))
    rho = asin(sin_rho)
    beta = asin(sin(_v3)/sin_rho)
    if _v2 < 0: beta = pi - beta
    gamma = asin(sin(_v2)/sin_rho)
    if _v3 < 0: gamma = pi - gamma
    A = pi/2. + _roll - beta
    B = atan2( sin(A)*cos(_dec), (sin(_dec)*sin_rho - cos(_dec)*cos(rho)*cos(A)))

    # compute final value
    troll = RADTODEG(pi - (gamma+B))

    return troll