from astropy.io import fits
import numpy as np
import pylenstool as plens
import os
import sys
import pandas as pd
from astropy.table import Table,Column
import subprocess

def mad(arr):
    """ Median Absolute Deviation: a "Robust" version of standard deviation.
        Indices variabililty of the sample.
        https://en.wikipedia.org/wiki/Median_absolute_deviation 
    """
    arr = np.ma.array(arr).compressed() # should be faster to not use masked arrays.
    med = np.median(arr)

    return np.median(np.abs(arr - med))

image_cat='../A2390_z_cleaned_20MAY21.mul'
bayesamp_file='magnification.cat' 
model="bestopt.par"
catalog="../A2390_z_cleaned_photom_20MAY21.fits"

liste_im_mul=[
        ('P307','1.1,1.2'),
        ('P315','3.1'),
        ('P312','3.2'),
        ('P67','3.3'),
        ('P206','4.1,4.2'),
        ('P318','4.3'),
        ('P632','4.3'),
        ('P309','5.1'),
        ('P308','5.2'),
        ('M104','6.1,6.2'),
        ('P150','6.3'),
        ('P280','7.1'),
        ('M114','7.2'),
        ('M109','9.1'),
        ('X1','9.2'),
        ('X2','9.3'),
        ('M131','10.1'),
        ('X3','10.2'),
        ('P140','11.1'),
        ('P156','11.2'),
        ('X4','11.3')]
        

im_mul=dict(liste_im_mul)

# add new columns to the catalog
t=Table.read(catalog)

c1=Column(name='MU',dtype='>f8',length=len(t))
c2=Column(name='MU_ERR',dtype='>f8',length=len(t))
c3=Column(name='MUL',dtype='S20',length=len(t))

t.add_column(c1)
t.add_column(c2)
t.add_column(c3)

# run bayesAmp
pf=plens.lenstool_param_file(model)
pf.set_parameter('runmode','inverse 0')
pf.set_parameter('runmode','image 0')
pf.set_parameter('image','multfile 1 '+str(image_cat))
pf.set_parameter('champ','dmax 30')

pf.write_param_file(output_filename=model)

os.system('bayesAmp -1 '+str(model)+' > '+str(bayesamp_file))

nlines=int(subprocess.check_output([""" grep -v '#' """+bayesamp_file+""" | wc -l"""],shell=True))

# read the mul file ans sort it by redshift
with open(image_cat) as file:
    datatable_im = pd.read_csv(image_cat,sep="\s+",header=None,skiprows=1)
datatable_im.sort_values([6],inplace=True)
datatable_im=datatable_im.reset_index(drop=True)
#print (datatable_im)

nb_sources=int(len(datatable_im)) # nb of sources

# read the file producted by bayesAmp, already sorted by redshift by lenstool
with open(bayesamp_file) as file:
    datatable_amp = pd.read_csv(bayesamp_file,sep="\s+",header=None,skip_blank_lines=True,skiprows=nb_sources+1)
    datatable_header=pd.read_csv(bayesamp_file,sep="\s+",skip_blank_lines=True,skipfooter=nlines,skiprows=1, header=None)

print (nb_sources)

for i in range(nb_sources):
    values=[]
    print (i)
    # get the values of mu for this ID
    for j in range(len(datatable_amp[i+1])):
        values.append(float(datatable_amp[i+1][j]))
    #print (values)
    source_name=str(datatable_header[2][i])
    ampl=np.mean(values) # measure mu
    ampl_err=1.68*mad(values) # measure mu_err
    print (source_name[1:-2])
    ksel=np.where((t['iden'][:]==float(source_name[1:-2]))) # find the object in the catalog
    #print (len(ksel[0]))
    print (ksel)
    for k in range(len(ksel[0])):
        if ((t['idfrom'][ksel[0][k]])=='PRIOR' and source_name[:1]=='P'):
            line = ksel[0][k]
            #print ('salut')
        if ((t['idfrom'][ksel[0][k]])=='MUSELET' and source_name[:1]=='M'):
            line = ksel[0][k]
            #print ('salu')
        if ((t['idfrom'][ksel[0][k]])=='EXTERN' and source_name[:1]=='X'):
            line = ksel[0][k]
            #print ('coucou')
    print (t['iden'][line])
    t['MU'][line]=ampl # add mu and mu_err
    t['MU_ERR'][line]=ampl_err
    # check if it is a multiple image or not 
    if str(source_name[:-2]) in im_mul:
        t['MUL'][line]=im_mul[str(source_name[:-2])]
    else :
        t['MUL'][line]=""
print ('coucou')
# put mu = 1 for all objects with no mu values
ksel=np.where(t['MU'][:]==0)
t['MU'][ksel]=1

t.write(catalog[:-5]+'_STEPB.fits',overwrite=True)