import sys
import numpy as np
import healpy as hp
from null_test import rotate
from scipy.stats import median_absolute_deviation as mad

sys.path.append('../sat_ephemeris')
from sat_ids import norad_ids

def good_maps(raw_tile):
    '''Here, we extract "good satellie" data from the tile_maps_raw 
    generated by tile_maps.py. We also normalize the maps by
    the reference beam'''

    f_name, _ = raw_tile.name.split('.')
    tile, ref, _, _ = f_name.split('_')
    
    # Load reference FEE model
    # Rotate the fee models by -pi/2 to move model from spherical (E=0) to Alt/Az (N=0)
    ref_fee_model = np.load(ref_model, allow_pickle=True)
    if 'XX' in tile:
        ref_fee = ref_fee_model['XX']
        rotated_fee = rotate(nside, angle=-(1*np.pi)/2.0,healpix_array=ref_fee)
    else:
        ref_fee = ref_fee_model['YY']
        rotated_fee = rotate(nside, angle=-(1*np.pi)/2.0,healpix_array=ref_fee)
   
    # load data from map .npz file
    tile_data   = np.load(raw_tile, allow_pickle=True)
    tile_data   = {key:tile_data[key].item() for key in tile_data}
    mwa_map     = tile_data['mwa_map']  
    
    tile_maps_norm = {p:[] for p in pointings}
    mwa_maps_norm = {p:[] for p in pointings}
     
    for p in pointings:
        
        # mwa map
        mwa_map_good = [[] for pixel in range(hp.nside2npix(nside))]
        
        for sat in good_sats:

            for pix in range(hp.nside2npix(nside)):

                mwa_map_good[pix].extend(mwa_map[p][sat][pix])

        mwa_maps_norm[p].extend(mwa_map_good)
        
    # Save map arrays to npz file
    np.savez_compressed(f'{out_dir}/{tile}_{ref}_tile_maps.npz', **mwa_maps_norm)


if __name__=='__main__':
    
    import argparse
    from pathlib import Path
    import concurrent.futures

    parser = argparse.ArgumentParser(description="""
        Extract the data from a list of 'good satellites' from the raw tile maps made my tile_maps.py. 
        Divide the tile/ref beams and apply the FEE ref model.
        """)
    
    parser.add_argument('--out_dir', metavar='\b', default='./../../outputs/tile_maps/tile_maps_norm/',
            help='Output directory. Default=./../../outputs/tile_maps/tile_maps_norm/')
    parser.add_argument('--map_dir', metavar='\b', default='./../../outputs/tile_maps/tile_maps_raw',
            help='Raw Map directory. Default=./../../outputs/tile_maps/tile_maps_raw')
    parser.add_argument('--ref_model', metavar='\b', default='../../outputs/reproject_ref/ref_dipole_models.npz',
            help='Healpix reference FEE model file. default=../../outputs/reproject_ref/ref_dipole_models.npz')
    parser.add_argument('--nside', metavar='\b', type=int,  default=32,help='Healpix Nside. Default = 32')
    
    args = parser.parse_args()
    
    out_dir     = Path(args.out_dir)
    map_dir     = Path(args.map_dir)
    ref_model   = Path(args.ref_model)
    nside       = args.nside
    
    # make outpud dir
    out_dir.mkdir(parents=True, exist_ok=True)
    
    # Good sats from which to make plots
    good_sats = [
            25338, 25982, 25984, 25985,
            28654, 40086, 40087,
            40091, 41179, 41180,
            41182, 41183, 41184,
            41185, 41187, 41188,
            41189, 44387
            ]

    # use all satellites
    #good_sats = list(norad_ids.values())

    # list of beam pointings
    pointings = ['0','2','4','41']
   
    # list of all raw tile maps
    map_files = [item for item in map_dir.glob('*.npz')]

    # Parallization magic happens here
    with concurrent.futures.ProcessPoolExecutor() as executor:
        results = executor.map(good_maps, map_files)

#    good_maps(map_files[0])






