#!/usr/bin/env python

import fitsio
import numpy as np
from compare_y1_y3 import fitting
from compare_y1_y3.util import (
    Y3_CHI_FACTORS,
    print_stats,
    add_rescaled_data,
    write_delta_sigma,
)

import dsfit


y3sources = '2pt_NG_final_2ptunblind_11_13_20_wnz.fits'

y3lenses_y3sources_y3area = '~/git/xcorr/runs/Y3_mastercat___UNBLIND___final_v1.0_DO_NOT_USE_FOR_2PT/zslim_som/zs_som/redmagic_x40randoms/zllim_y3/lens_w_True/njk_150/thbin_2.50_250_20/bslop_0/source_only_close_to_lens_True_nside4/measurement/gt_boosted_twopointfile.fits'  # noqa
y3lenses_y3sources_y1area = '~/git/xcorr/runs/Y3_mastercat___UNBLIND___final_v1.0_DO_NOT_USE_FOR_2PT_y1footprint/zslim_som/zs_som/redmagic_x40randoms_year1footprint/zllim_y3/lens_w_True/njk_150/thbin_2.50_250_20/bslop_0/source_only_close_to_lens_True_nside4/measurement/gt_boosted_twopointfile.fits'  # noqa

# (sbin, lbin),
BINS2FIT = [
    (1, 1),
    (1, 2),
    (1, 3),  # marginal
    (1, 4),  # *
    # (1, 5),  # *
    #
    (2, 1),
    (2, 2),
    (2, 3),
    (2, 4),   # *
    # (2, 5),   # *
    #
    (3, 1),
    (3, 2),
    (3, 3),
    (3, 4),  # *
    # (3, 5),  # *
    #
    (4, 1),
    (4, 2),
    (4, 3),
    (4, 4),  # *
    # (4, 5),  # *

]


def get_args():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', type=int, required=True)
    return parser.parse_args()


def read_data():
    data = {
        'y1area': {
            'gammat': fitsio.read(y3lenses_y3sources_y1area, ext='gammat', lower=True),  # noqa
            'gammat_cov': fitsio.read(y3lenses_y3sources_y1area, ext='covmat'),  # noqa
            # 'nzl': fitsio.read(y3lenses_y3sources_y1area, ext='nz_lens', lower=True),  # noqa
            # 'nzs': fitsio.read(y3lenses_y3sources_y1area, ext='nz_source', lower=True),  # noqa
            'nzl': fitsio.read(y3sources, ext='nz_lens', lower=True),  # noqa
            'nzs': fitsio.read(y3sources, ext='nz_source', lower=True),  # noqa
            'chi_factors': Y3_CHI_FACTORS,
        },
        'y3area': {
            'gammat':  fitsio.read(y3lenses_y3sources_y3area, ext='gammat', lower=True),  # noqa
            'gammat_cov':  fitsio.read(y3lenses_y3sources_y3area, ext='covmat'),  # noqa
            # 'nzl': fitsio.read(y3lenses_y3sources_y3area, ext='nz_lens', lower=True),  # noqa
            # 'nzs': fitsio.read(y3lenses_y3sources_y3area, ext='nz_source', lower=True),  # noqa
            'nzl': fitsio.read(y3sources, ext='nz_lens', lower=True),  # noqa
            'nzs': fitsio.read(y3sources, ext='nz_source', lower=True),  # noqa
            'chi_factors': Y3_CHI_FACTORS,
        }
    }

    return data


def do_fits(*, rng, data, cosmo):

    tab_y1area = fitting.make_table()
    tab_y3area = fitting.make_table()
    y1area_reslist = []
    y3area_reslist = []

    iax = 0
    for sbin, lbin in BINS2FIT:
        res_y1area = fitting.fit_nfw_lin(
            rng=rng, data=data['y1area'],
            lbin=lbin, sbin=sbin, plt=tab_y1area.axes[iax],
            cosmo=cosmo,
        )
        res_y3area = fitting.fit_nfw_lin(
            rng=rng, data=data['y3area'],
            lbin=lbin, sbin=sbin, plt=tab_y3area.axes[iax],
            cosmo=cosmo,
        )

        y1area_reslist.append(res_y1area)
        y3area_reslist.append(res_y3area)

        iax += 1

    tab_y1area.axes[0].legend()
    tab_y3area.axes[0].legend()

    return y1area_reslist, y3area_reslist, tab_y1area, tab_y3area


def main():
    args = get_args()
    rng = np.random.RandomState(args.seed)

    cosmo = dsfit.get_cosmo('planck18')

    data = read_data()

    add_rescaled_data(
        data=data['y1area'], source_type='y3',
        cosmo=cosmo,
    )
    add_rescaled_data(
        data=data['y3area'], source_type='y3',
        cosmo=cosmo,
    )

    y1area_ds_file = 'y3lenses-y3sources-y1area-delta-sigma.fits'
    write_delta_sigma(filename=y1area_ds_file, data=data['y1area'])
    y3area_ds_file = 'y3lenses-y3sources-y3area-delta-sigma.fits'
    write_delta_sigma(filename=y3area_ds_file, data=data['y3area'])

    y1area_reslist, y3area_reslist, tab_y1area, tab_y3area = do_fits(
        rng=rng, data=data, cosmo=cosmo,
    )
    tab = print_stats(
        reslist1=y1area_reslist, reslist2=y3area_reslist,
        label1='y1 area', label2='y3 area',
        title='y3 lenses and sources',
    )
    pngname = 'y3lenses-y3sources-diffarea-weighthist.png'
    print('writing:', pngname)
    tab.savefig(pngname, dpi=100)

    dpi = 75

    y1area_fname = 'y3lenses-y3sources-y1area-fits.png'
    y3area_fname = 'y3lenses-y3sources-y3area-fits.png'

    print('writing:', y1area_fname)
    tab_y1area.savefig(y1area_fname, dpi=dpi)

    print('writing:', y3area_fname)
    tab_y3area.savefig(y3area_fname, dpi=dpi)

    # print('amp: %g +/- %g' % (amp, amp_err))


if __name__ == '__main__':
    main()
