"""**Obsolete** We are using scale factor PSF fitting instead.

Renormalized FHD peak-normalize maps to integral-normalize maps by
performing polynomial fitting with Gaussian maps.

"""
import argparse
from multiprocessing import Pool, Manager

import matplotlib
import numpy as np

matplotlib.use('Agg')
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.modeling import models, fitting
from astropy.time import Time

import opstats.utils.settings as s
from opstats.utils import check_dir


_mgr = Manager()
_config = _mgr.Namespace()
_config.root_dir = '/data3/piyanat/runs/eor/ha0000h/uvlt50/'
_config.maps_dir = '/data3/piyanat/runs/eor/ha0000h/uvlt50/maps/'
_config.res_dir = '/data3/piyanat/runs/eor/ha0000h/uvlt50/res/'
_config.freq = s.MWA.FREQ_EOR_ALL_80KHZ
_res_scale = _mgr.Array('d', np.zeros_like(_config.freq))


def fit_residual(ch):
    f = _config.freq[ch]
    fhd_map = fits.open(_config.maps_dir + 'fhd_0.000h_{:.3f}MHz.fits'
                        .format(f))
    smooth_map = fits.open(_config.maps_dir + 'gauss_0.000h_{:.3f}MHz.fits'
                           .format(f))
    mask = ~np.isnan(fhd_map[0].data)
    x = fhd_map[0].data[mask]
    y = smooth_map[0].data[mask]
    # Fit pixel value
    p_init = models.Polynomial1D(1)
    p_init.fixed['c0'] = True
    fitter = fitting.LinearLSQFitter()
    p = fitter(p_init, x, y)
    scale_factor = p.c1.value
    smooth_map[0].data /= scale_factor
    res_map = fhd_map[0].data - smooth_map[0].data
    smooth_map[0].header['DATE'] = (Time.now().isot, 'Date of file creation')
    smooth_map.writeto(_config.maps_dir + 'gauss_scaled_0.000h_{:.3f}MHz.fits'
                       .format(f), clobber=True)
    hdu = fits.PrimaryHDU(data=res_map, header=smooth_map[0].header)
    hdu.header['DATE'] = (Time.now().isot, 'Date of file creation') 
    hdu.writeto(_config.maps_dir + 'res_0.000h_{:.3f}MHz.fits'.format(f),
                clobber=True)
    # save scaling to shared array
    _res_scale[ch] = scale_factor
    fhd_map.close()
    smooth_map.close()
    # Make a plot
    plt.plot(x, y, 'b.', label='data pair', markersize=1)
    xi = np.linspace(x.min(), x.max())
    plt.plot(xi, p(xi), 'r-', label='fit: slope = {:.3f}'.format(p.c1.value))
    plt.xlabel('Gaussian Smoothed Map Pixel Values (mK)')
    plt.ylabel('FHD Map Pixel Values (mK)')
    plt.legend(loc='upper left', numpoints=1)
    plt.title('Residual Fit {:.3f} MHz'.format(f))
    plt.xlim(-500, 500)
    plt.ylim(-6000, 6000)
    plt.tight_layout()
    plt.minorticks_on()
    plt.savefig(_config.res_dir + 'res_fit_plot_0.000h_{:.3f}MHz.png'.format(f))
    plt.close()


def main(root_dir='/data3/piyanat/runs/eor/ha0000h/uvlt50/', nprocs=16):
    _config.root_dir = root_dir
    _config.res_dir = root_dir + '/res/'
    _config.maps_dir = root_dir + '/maps/'
    check_dir(_config.res_dir)
    check_dir(_config.maps_dir)
    pool = Pool(nprocs)
    pool.map(fit_residual, range(_config.freq.size))
    pool.close()
    pool.join()
    res_scale = np.array(_res_scale)
    np.savetxt(_config.res_dir + 'res_scale_0.000h.csv', res_scale,
               delimiter=',')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_dir', type=str,
                        default='/data3/piyanat/runs/eor/ha0000h/uvlt50/',
                        help='root directory of simulation data')
    parser.add_argument('--nprocs', type=int, default=16,
                        help='number of processes to spawn')
    args = parser.parse_args()
    main(root_dir=args.root_dir, nprocs=args.nprocs)
