import healpy as hp
import numpy as np
from astropy.wcs import WCS

class SinProjection:
    def __init__(self, brick, pixel_scale=0.074):

        self.brick = brick
        self.cdelt = pixel_scale / 3600.0  
        
        self.wcs = WCS(naxis=2)
        self.wcs.wcs.crval = [self.brick.ra, self.brick.dec]
        self.wcs.wcs.crpix = [20000, 20000]
        self.wcs.wcs.cdelt = [-self.cdelt, self.cdelt]
        self.wcs.wcs.ctype = ["RA---SIN", "DEC--SIN"]

    def healpix_original_boundary(self):
        """
        Get the original HEALPix pixel boundary in RA and Dec.

        Returns
        -------
        tuple of ndarray
            ra_v : ndarray
                Array of vertex Right Ascension values (degrees).
            dec_v : ndarray
                Array of vertex Declination values (degrees).
        """
        nside = self.brick.nside 
        verts = hp.boundaries(nside, self.brick.brick_id, step=1)
        ra_v = np.degrees(np.arctan2(verts[1], verts[0])) % 360
        dec_v = np.degrees(np.arcsin(verts[2]))
        return ra_v, dec_v
    

    def expanded_brick_boundary(self, n_points=60):
        """
        Generate an expanded circular boundary for the brick based on its radius.

        Parameters
        ----------
        n_points : int, optional
            Number of points used to sample the circular boundary. Default is 60.

        Returns
        -------
        tuple of list
            ras : list of float
                List of sampled Right Ascension values (degrees).
            decs : list of float
                List of sampled Declination values (degrees).
        """
        ra_c, dec_c = np.radians(self.brick.ra), np.radians(self.brick.dec)
        radius = np.radians(self.brick.radius)
        ras, decs = [], []
        for theta in np.linspace(0, 2*np.pi, n_points):
            dec = np.arcsin(np.sin(dec_c)*np.cos(radius) +
                            np.cos(dec_c)*np.sin(radius)*np.cos(theta))
            ra = ra_c + np.arctan2(np.sin(theta)*np.sin(radius)*np.cos(dec_c),
                                   np.cos(radius)-np.sin(dec_c)*np.sin(dec))
            ras.append(np.degrees(ra) % 360)
            decs.append(np.degrees(dec))
        return ras, decs

    def project_to_sin(self, mode="brick", n_points=60):
        """
        Generate RA and Dec boundary points for a brick or HEALPix pixel.

        Parameters
        ----------
        mode : {'brick', 'healpix'}, optional
            - 'brick' : Use the expanded brick boundary (`expanded_brick_boundary`).
            - 'healpix' : Use the original HEALPix pixel boundary (`healpix_original_boundary`).
            Default is 'brick'.
        n_points : int, optional
            Number of points to sample along the circular boundary. Default is 60.

        Returns
        -------
        tuple of list
            ra_v : list of float
                List of Right Ascension values (degrees).
            dec_v : list of float
                List of Declination values (degrees).

        Raises
        ------
        ValueError
            If `mode` is not 'brick' or 'healpix'.
        """
        if mode == "brick":
            ra_v, dec_v = self.expanded_brick_boundary(n_points=n_points)
        elif mode == "healpix":
            ra_v, dec_v = self.healpix_original_boundary()
        else:
            raise ValueError("mode must be either 'brick' or 'healpix'")

        return ra_v, dec_v 
    
    
    def world2pix(self, ra_v, dec_v):
        """
        Convert RA and Dec coordinates to pixel (X, Y) coordinates 
        using the current WCS (SIN projection).

        Parameters
        ----------
        ra_v : array_like
            Right Ascension values (degrees).
        dec_v : array_like
            Declination values (degrees).

        Returns
        -------
        tuple of ndarray
            x : ndarray
                X pixel coordinates corresponding to the input RA/Dec.
            y : ndarray
                Y pixel coordinates corresponding to the input RA/Dec.
        """
        return self.wcs.wcs_world2pix(ra_v, dec_v, 0)