#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 27 21:23:09 2023

corr_ref_pxs_simple3.py

Se basa en corr_ref_pxs_simple2, con más opciones

Se resta a la rampa / cds el valor del primer frame que se vaya a usar.


Oct24: añado tratamiento de resets y posibilidad de eliminarlo en
el cálculo.

@author: fgl
"""

import numpy as np
import sys
import os
from astropy.io import fits
from datetime import datetime

frmt=1.47579
lsegm=64   # tamaño del segmento horizontal
nsegm=32    # no. de segmentos 
ybox=4  # tamaño de la caja para mediana
grad=15    # grado del polinomio de ajuste a las refs verticales
ncols=2048
nref=4

def linfit_2d_gral(x,y):
    # sirve para puntos no equiespaciados, e.g, CDS/FOWLER
    # 1er punto de ajuste es 1. Si usamos el dt -> adu/s
    nfr=len(x)
    sx=np.sum(x)   
    fac = nfr*np.sum(x*x)-sx*sx
    slope=(nfr*np.sum(x*y,2)-sx*np.sum(y,2))/fac
    return slope

def linfit_2d(x,y):
    # sirve para puntos equiespaciados, e.g, RAMP
    # 1er punto de ajuste es 1. No usamos el dt -> adu/dx
    nfr=len(x)
    fac = 12 / (nfr*nfr - 1)/ nfr
    slope=(np.sum(x*y,2)-(nfr+1)*np.sum(y,2)/2)*fac
    return slope

nargv=len(sys.argv) - 1
if nargv<2:
    print('This script needs 2 arguments, at least: ')
    print('\t* Input folder')
    print('\t* Output folder')
    print('\t* Corr type: Href - 0; HVref - 1, default 0')
    print('\t* ACN correction, 1 - y, 0 - n, default 0')
    print('\t* Remove reset if used, 1 - y, 0 - n, default 0')
    sys.exit()

dir_in=sys.argv[1]
if dir_in[-1] != '/':
    dir_in += '/'

dir_out=sys.argv[2]
if dir_out[-1] != '/':
    dir_out += '/'
if not os.path.exists(dir_out):
    os.makedirs(dir_out)

if nargv > 2:
    vcor=int(sys.argv[3])
else:
    vcor=False

if nargv > 3:
    ACN=int(sys.argv[4])
else:
    ACN=False

if nargv > 4:
    RST=int(sys.argv[5])
else:
    RST=False
    
filelist = [dir_in+f for f in os.listdir(dir_in) \
            if os.path.isfile(os.path.join(dir_in, f)) \
            if '_raw'  in f and '.fits' in f]          

filelist.sort()
nfich=len(filelist)
nf=0
# xpp=np.arange(1,2041) # no usamos los pxs de referencia
xpp=np.arange(nref-ncols//2,ncols//2-nref) # no usamos los pxs de referencia
while nf < nfich:
    file=filelist[nf]
    grp=int(fits.getval(file,'GROUPNUM'))
    rd=int(fits.getval(file,'READNUM'))
    # print ('\t\t\t',file.split('/')[-1],grp,rd)
    # if itern0 !=1:
    ngroups=int(fits.getval(file,'NGROUPS'))
    nreads=int(fits.getval(file,'NREADS'))
    nrsts=int(fits.getval(file,'NRESETS'))
    try:
        rstn=int(fits.getval(file,'resetnum'))
    except:
        rstn=0
    if (grp*rd !=1 and rstn == 0) or rstn > 1:  # no es primer fichero
        nf += 1
        continue
    nimgt=ngroups*nreads+(rstn>0)
    if nimgt == 1:
        nf += 1
        continue
    if rstn != 0:  # buscamos el último reset
        while rstn < nrsts:
            nf += 1
            file=filelist[nf]
            rstn=int(fits.getval(file,'resetnum'))
            
    # texp0=float(fits.getval(file,'EXPTIME'))
    suff='-'+('-').join(file.split('-')[1:])    
    tfow=float(fits.getval(file,'TFOWLER'))/1000 # ya en segundos
    
    texp=(nimgt-1)*frmt+tfow
    # if abs(texp-texp0)> 1e-3:
    #     print ('\t\t\t****Texp wrong in header, ',str('%.3f'%texp0), \
    #            str('%.3f'%texp),str(ngroups),str('%.3f'%tfow))

    cube=np.zeros((ncols,ncols,nimgt))
    nimg=0
    cube[:,:,nimg]=fits.getdata(file).astype(float)
    print (file.split('/')[-1],nrsts,ngroups,nreads,rstn,grp,rd)
    nimg=1
    while nimg < nimgt:
        nf += 1
        file=filelist[nf]
        grp=int(fits.getval(file,'GROUPNUM'))
        assert grp == nimg+(rstn==0), 'Error in ramp. ' + \
            ' Check completness of ist of files'
        cube[:,:,nimg]=fits.getdata(file).astype(float)
        print ('\t',file.split('/')[-1],grp)
        nimg += 1
    
    if RST and rstn > 0:
        nrv=1
        nimgt -= 1
        texp -= frmt
        cube = np.delete(cube,0,-1)
    cube[:,:] -= np.stack([cube[:,:,0]]*nimgt,axis=-1)
    # for i in range(nimgt):
    #     cube[:,:,i] -= cube[:,:,0]
        
    id=int(file.split('/')[-1].split('-')[0])+1
    suff=file.split('/')[-1].split('-')[1:]
    suff='-'.join(suff)[:-9]+'.fits'
    # suff=str('-'.join(suff).split('_')[:-1])+'.fits'
    file=str("%010d" % id)+'-'+suff
    if not os.path.exists(dir_in+file):
        nf += 1
        print ('\t Preprocessed file missing. Skipping.')
        continue
        
    refb=np.median(cube[:nref],0)
    reft=np.median(cube[-nref:],0)
    if ACN:
        lsegm2 = lsegm/2
        # tth=np.zeros((nsegm*2,nimgt))
        tth_e=np.array([np.mean((refb[i*lsegm:i*lsegm+lsegm:2]+ \
            reft[i*lsegm:i*lsegm+lsegm:2])/2, \
                    axis=0) for i in range(nsegm)])
        tth_o=np.array([np.mean((refb[i*lsegm+1:i*lsegm+lsegm:2]+ \
            reft[i*lsegm+1:i*lsegm+lsegm:2])/2, \
                    axis=0) for i in range(nsegm)])
    else:
        tth=np.array([np.mean((refb[i*lsegm:i*lsegm+lsegm]+ \
            reft[i*lsegm:i*lsegm+lsegm])/2, \
                    axis=0) for i in range(nsegm)])

    cube1=np.copy(cube)
    if vcor:
        # refl=np.median(cube[:,:nref],1)
        # refr=np.median(cube[:,-nref:],1)
        # ttv=np.array([(np.median(refl[i-ybox:i+ybox+1],axis=0) +\
        #             np.median(refr[i-ybox:i+ybox+1],axis=0))/2 \
        #             for i in range(ybox,ncols-ybox)])
        ppl=np.polyval(np.polyfit(xpp,np.median(cube[nref:-nref,:nref],1), \
                      grad), np.array([xpp]*nimgt).T)
        ppr=np.polyval(np.polyfit(xpp,np.median(cube[nref:-nref,:nref],1), \
                      grad), np.array([xpp]*nimgt).T)
        ttv=(ppl+ppr)/2
        ttvnorm=ttv/np.mean(ttv,0)
        for i in range(nimgt):
          if ACN:
            cube1[ybox:-ybox,::2,i]=cube[ybox:-ybox,::2,i] - \
                    np.repeat(np.outer(ttvnorm[:,i],tth_e[:,i]),lsegm2,1)
            cube1[ybox:-ybox,1::2,i]=cube[ybox:-ybox,1::2,i] - \
                    np.repeat(np.outer(ttvnorm[:,i],tth_o[:,i]),lsegm2,1)
          else:
            cube1[ybox:-ybox,:,i]=cube[ybox:-ybox,:,i] - \
                    np.repeat(np.outer(ttvnorm[:,i],tth[:,i]),lsegm,1)
    else:
        if ACN:
          cube1[:,::2]=cube[:,::2]-np.tile(np.repeat(tth_e,lsegm2,1), \
                                     (ncols,1,1))
          cube1[:,1::2]=cube[:,1::2]-np.tile(np.repeat(tth_o,lsegm2,1), \
                                             (ncols,1,1))
        else:
          cube1=cube - np.tile(np.repeat(tth,lsegm,0),(ncols,1,1))
        
    if nimgt==0:
        slope=cube1[:,:,1]
    elif tfow==0:
        slope=linfit_2d(np.arange(1,nimgt+1),cube1)*texp/frmt
    else:
        xp=np.arange(1,nimgt+1).astype(float)
        xp[nreads:] += tfow/frmt
        xp -= xp.mean()
        slope=linfit_2d_gral(xp,cube1)*texp/frmt
    
            
    head1=fits.getheader(dir_in+file,ext=0)
    head2=fits.getheader(dir_in+file,ext=1)
    hduw = fits.PrimaryHDU(slope.astype('float32'))
    head1['bzero']=0
    head1['exptime']=texp
    
    date=datetime.today().strftime('%Y-%m-%d')   
    head1['history']=date + ' Horizontal refs. pixel subtracted'
    if vcor and RST+ACN==0:
        head1['history']=date + ' Vertical refs. pixel used for ' + \
                        'polynomial fitting of ref. variation'
    if  ACN:
        head1['history']=date + ' Alternate columns noise used in ' + \
                    'subtraction'
    if  RST:
        head1['history']=date + ' Reset frame eliminated in the processing'

    # hduw.scale('int16', bzero=32768,bscale=1)
    hduw.header=head1
    hduw2=fits.ImageHDU(header=head2)
    hdul = fits.HDUList([hduw, hduw2])
    hdul.writeto(dir_out+file, overwrite=True)   
    # print('\tWritten ',dir_out+file)
    print('\n\t\tWritten ',file,'\n')
    nf += 1