from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status

from .models import GaiaSource
from .serializers import GaiaSourceSerializer

import numpy as np
import healpy as hp
from astropy.coordinates import SkyCoord
from astropy import units as u


class ConeSearchView(APIView):

    @staticmethod
    def get_healpix_indices(ra, dec, radius, nside):

        # Convert the input coordinates to a SkyCoord object
        center = SkyCoord(ra=ra*u.deg, dec=dec*u.deg, frame='icrs')

        # Convert the center coordinates to HEALPix vector
        vec = hp.ang2vec(center.ra.deg, center.dec.deg, lonlat=True) 

        # Find the pixels within the given radius
        indices = hp.query_disc(nside, vec, np.radians(radius))        
        return indices

    def get(self, request, ra, dec, radius):
        radius = float(radius)
        ra = float(ra)
        dec = float(dec)
        nside = 2048

        #create skycoord for the center of search
        center = SkyCoord(ra=ra*u.deg, dec=dec*u.deg, frame='icrs')

        #fetch healpix indices in the specified disc
        healpix_indices = self.get_healpix_indices(ra, dec, radius, nside)
        #fetch all objects from those healpixes
        objects = GaiaSource.objects.filter(healpix_ring_index__in=healpix_indices)

        results = [] #initialize the results list

        for obj in objects:
            
            source_coordinates = SkyCoord(
                ra=obj.ra*u.deg, 
                dec=obj.dec*u.deg, 
                frame='icrs'
            )
            
            separation = center.separation(source_coordinates)

            if separation.degree <= radius:

                results.append(obj)           

        serializer = GaiaSourceSerializer(results, many=True)

        return Response(serializer.data, status=status.HTTP_200_OK)