# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2019-2020 Max-Planck-Society
# Author: Philipp Arras

import matplotlib.pyplot as plt
import numpy as np

import nifty6 as ift
import src as vlbi
from config import comm, rank, master
from config import oversampling_factor as ofac
from config import sky_movie_mf as sky



def main():

    # Prior plots
    if master:
        with ift.random.Context(31):
            p = ift.Plot()
            n = 5
            for _ in range(5):  # FIXME: should this be 'range(n)'?
                pos = ift.from_random(sky.domain, 'normal')
                ss = sky(pos)['hi']
                mi, ma = 0, np.max(ss.val)
                for ii in range(0, ss.shape[0], 4):
                    extr = ift.DomainTupleFieldInserter(ss.domain, 0, (ii,)).adjoint
                    p.add(extr(ss), zmin=mi, zmax=ma)
            p.output(name=f'prior_samples.png', ny=n, xsize=28, ysize=9)

            print('Start inf check')
            for ii in range(20):
                pp = ift.from_random(sky.domain, 'normal') #* 2.5
                sky(pp)

    fld = vlbi.gaussian_profile(sky.target['hi'][1], 30 *vlbi.MUAS2RAD)
    fld = ift.ContractionOperator(sky.target['hi'], 0).adjoint(fld).val
    multi = ift.makeField(sky.target, {'lo': fld, 'hi': fld})
    output = vlbi.normalize(multi.domain)(multi)

    cov = 1e-3*max([vv.max() for vv in output.val.values()])**2
    dtype = list(set([ff.dtype for ff in output.values()]))
    if len(dtype) != 1:
        raise ValueError('Only MultiFields with one dtype supported.')
    dtype = dtype[0]
    invcov = ift.ScalingOperator(output.domain, cov).inverse
    d = output + invcov.draw_sample_with_dtype(dtype, from_inverse=True)
    lh = ift.GaussianEnergy(d, invcov) @ sky
    H = ift.StandardHamiltonian(
        lh, ic_samp=ift.AbsDeltaEnergyController(0.5, iteration_limit=100, convergence_level=2, name=f'CG(task {rank})'))
    pos = 0.1*ift.from_random(sky.domain, 'normal')
    cst = ('spaceasperity', 'spaceflexibility', 'spacefluctuations', 'spaceloglogavgslope', 'spacespectrum', 'timeasperity', 'timeflexibility', 'timefluctuations', 'timeloglogavgslope', 'timespectrum')
    minimizer = ift.NewtonCG(ift.GradientNormController(iteration_limit=10, name='findpos' if master else None))
    n = 2
    for ii in range(n):
        if master:
            ift.logger.info(f'Start iteration {ii+1}/{n}')
        kl = ift.MetricGaussianKL(pos, H, 2, comm=comm, mirror_samples=True, constants=cst, point_estimates=cst)
        kl, _ = minimizer(kl)
        pos = kl.position

        if master:
            movie = vlbi.freq_avg(vlbi.strip_oversampling(sky(pos), ofac)).val
            ma = np.max(movie)
            fig, axs = plt.subplots(5, 6, figsize=(12, 12))
            axs = axs.ravel()
            for ii in range(movie.shape[0]):
                axx = axs[ii]
                axx.imshow(movie[ii].T, vmin=0, vmax=ma)
                axx.set_title(f'Frame {ii}')
            fig.savefig(f'movie_start.png')
            plt.close()

    if master:
        vlbi.save_hdf5('initial.h5', pos)
        with open("initial_random_state.txt", "wb") as f:
            f.write(ift.random.getState())


if __name__ == '__main__':
    # Change seed for the whole reconstruction here
    # ift.random.push_sseq_from_seed(42)
    main()