from __future__ import division
from pyoperators import pcg, DiagonalOperator, UnpackOperator
from pysimulators.noises import _gaussian_psd_1f
from qubic import (
    QubicAcquisition, create_sweeping_pointings, equ2gal)
import healpy as hp
import matplotlib.pyplot as mp
import numpy as np
import qubic

# read the input map
x0 = qubic.io.read_map(qubic.data.PATH + 'syn256_pol.fits', field='I_STOKES')

# instrument model
sigma = 10
fknee = 1
fslope = 1
ncorr = 10

# observation model
np.random.seed(0)
racenter = 0.0
deccenter = -57.0
angspeed = 1  # deg/sec
delta_az = 15.
angspeed_psi = 0.1
maxpsi = 45.
nsweeps_el = 300
duration = 24   # hours
ts = 20         # seconds
sampling = create_sweeping_pointings(
    [racenter, deccenter], duration, ts, angspeed, delta_az, nsweeps_el,
    angspeed_psi, maxpsi)

# acquisition model
acq = QubicAcquisition(150, sampling, kind='I', synthbeam_fraction=0.99,
                       detector_sigma=sigma, detector_fknee=fknee,
                       detector_fslope=fslope, detector_ncorr=ncorr)
C = acq.get_convolution_peak_operator()
P = acq.get_projection_operator()
H = P * C

# produce the Time-Ordered data
y = H(x0)

# noise
psd = _gaussian_psd_1f(len(acq.sampling), sigma=sigma, fknee=fknee,
                       fslope=fslope, sampling_frequency=1/ts)
invntt = acq.get_invntt_operator()
noise = acq.get_noise()

# map-making
coverage = P.pT1()
mask = coverage > 10
P = P.restrict(mask, inplace=True)
unpack = UnpackOperator(mask)

# map without covariance matrix
solution1 = pcg(P.T * P, P.T(y + noise),
                M=DiagonalOperator(1/coverage[mask]), disp=True)
x1 = unpack(solution1['x'])

# map with covariance matrix
solution2 = pcg(P.T * invntt * P, (P.T * invntt)(y + noise),
                M=DiagonalOperator(1/coverage[mask]), disp=True)
x2 = unpack(solution2['x'])


def display(x, title):
    x = x.copy()
    x[~mask] = np.nan
    hp.gnomview(x, rot=center, reso=5, xsize=600, min=-200, max=200,
                title=title)

center = equ2gal(racenter, deccenter)
display(C(x0), 'Original convolved map')
#hp.projplot(np.radians(pointings[..., 0]), np.radians(pointings[..., 1]))
display(x1, 'Reconstructed map no invntt')
display(x2, 'Reconstructed map with invntt')
display(coverage, 'Coverage')

mp.show()
