import glob
from astropy.io import fits
import numpy as np
import pandas as pd
import montage_wrapper as montage
import os.path
import pdb


input_folder = "./INPUT_FOLDER/"
output_folder = "./OUTPUT_FOLDER/"
vis_path = "./"
vis_image = "EUC_VIS_SWL-STK-000-000000-0000000__20210613T225647.742318Z.fits"


#I open the VIS image
hdu = fits.open(vis_path+vis_image)
img = hdu[1].data
hdr = hdu[1].header
dim1 = hdr["NAXIS1"]
dim2 = hdr["NAXIS2"]

if not os.path.exists(output_folder):
    os.mkdir(output_folder)

#these are the only files I will accept as inputs
filename_sims =  glob.glob(input_folder+"galaxy_and_stream_convolved_[!w]*.fits") #in the termina use ^ instead of !

for filename_sim in filename_sims:
    print(filename_sim)
    #I open each of the simulated galaxies
    hdu_sim = fits.open(filename_sim)
    img_sim = hdu_sim[1].data
    hdr_sim = hdu_sim[1].header
    dim1_sim = hdr_sim["NAXIS1"]
    dim2_sim = hdr_sim["NAXIS2"]
    
    flag_cutted = False
    while flag_cutted == False:
        #I obtain the new random centroid within the VIS image
        x_center = np.random.randint(0, dim1)
        y_center = np.random.randint(0, dim2)
        
        #I will take a cutout in the new centroid position with the dimensions of the simulated galaxy
        pix_in_x_halfsize = int(dim1_sim/2)
        pix_in_y_halfsize = int(dim2_sim/2)
        cutout = img[y_center-pix_in_y_halfsize:y_center+pix_in_y_halfsize, x_center-pix_in_x_halfsize:x_center+pix_in_x_halfsize]
        
        #I will accept this new image if less than 10% of the pixels are zeros and the dimensions are all right
        if np.count_nonzero(cutout == 0) >= (dim1_sim*dim2_sim)/10 or cutout.shape[0] != dim1_sim or cutout.shape[1] != dim2_sim:
            flag_cutted = False
        else:
            flag_cutted = True
            
    #I create the cutout itself
    new_filename = output_folder+filename_sim[len(input_folder):-5]+"_in_"+vis_image[:-5]+".fits"
    montage.mSubimage_pix(vis_path+vis_image,new_filename, x_center-pix_in_x_halfsize, y_center-pix_in_y_halfsize, hdu = 1, xpixsize = (pix_in_x_halfsize*2)-1, ypixsize = (pix_in_y_halfsize*2)-1 )
    hdu_cutout = fits.open(new_filename)
    img_cutout = hdu_cutout[0].data
    hdr_cutout = hdu_cutout[0].header
    fits.writeto(new_filename,img_cutout+img_sim,hdr_cutout,overwrite=True)