import numpy as np
from astropy.io import fits

from django.core.management.base import BaseCommand
from django.db import transaction
from axc_ul.models import Pixel

from itertools import islice

BATCH_SIZE = 1000000

def batch(iterable, size):
    """
    Generator that yields successive chunks of size 'size' from 'iterable'.
    """
    iterable = iter(iterable)
    while True:
        chunk = list(islice(iterable, size))
        if not chunk:
            break
        yield chunk


class Command(BaseCommand):
    help = "Process FITS files and store the data in the database"

    def add_arguments(self, parser):
        parser.add_argument(
            '--counts', 
            type=str, 
            required=True,
            help='Path of the counts file'
        )
        parser.add_argument(
            '--exposure', 
            type=str, 
            required=True,
            help='Path of the exposure file'
        )
        # parser.add_argument(
        #     '--survey_number', 
        #     type=int, 
        #     required=True,
        #     help='Integer ID of the survey being read'
        # )

    def handle(self, *args, **options):
        
        counts_file = options['counts']
        exposure_file = options['exposure']
        # survey_number = options['survey_number']

        self.stdout.write(f"Counts file:   {counts_file}")
        self.stdout.write(f"Exposure file: {exposure_file}")

        with fits.open(counts_file) as hdul:

            column_name = "T"
            counts_map = hdul[1].data[column_name]

        counts_data = counts_map.ravel()


        with fits.open(exposure_file) as hdul:

            column_name = "T"
            exposure_map = hdul[1].data[column_name]

        exposure_data = exposure_map.ravel()



        self.stdout.write(f"Counts Data Shape:   {counts_data.shape}")
        self.stdout.write(f"Exposure Data Shape: {exposure_data.shape}")

        assert counts_data.shape == exposure_data.shape, "Counts and exposure maps must have the same shape"

        #rate_data = np.divide(counts_data, exposure_data)
        
        # with transaction.atomic():
            
        #     survey,created = Survey.objects.get_or_create(number=survey_number)
            
        #     if created:
        #         self.stdout.write(f"Created a new survey instance with number: {survey.number}")
        #     else:
        #         self.stdout.write(f"Using existing survey instance with the number: {survey.number}")
        

        # Create a generator that yields Pixel objects one by one.
        pixel_generator = (
            Pixel(
                hpid=i,
                counts=int(count),
                exposure=float(exposure),
                #rate=float(rate),
                #survey=survey
            )
            for i, (count, exposure) in enumerate(zip(counts_data, exposure_data))
        )

        total_inserted = 0
        # Process the generator in batches.
        for pixel_batch in batch(pixel_generator, BATCH_SIZE):
            with transaction.atomic():
                Pixel.objects.bulk_create(pixel_batch)
            total_inserted += len(pixel_batch)
            self.stdout.write(f"Inserted {total_inserted} pixels")

        self.stdout.write(f"Inserted a total of {total_inserted} pixels.")