# Copyright (C) 2024 Richard Stiskalek
# 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, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
"""Mock data generators."""
import numpy as np

from ..field.interp import evaluate_cartesian_regular
from ..params import SPEED_OF_LIGHT
from ..utils import cartesian_to_radec, radec_to_cartesian, radec_to_galactic
from .cosmography import distmod2dist, distmod2redshift

###############################################################################
#                        Mock Quijote observations                            #
###############################################################################


def mock_Carrick2MTF(velocity_field, boxsize, RA_2MTF, DEC_2MTF,
                     a_TF=-22.8, b_TF=-7.2, sigma_TF=0.25, sigma_v=100.,
                     Vext=[150.0, 10.0, -100.0], beta=0.4,
                     mean_eta=0.069, std_eta=0.078, mean_e_eta=0.012,
                     mean_mag=10.31, std_mag=0.83, mean_e_mag=0.044,
                     seed=42, Om0=0.3, **kwargs):
    """
    Mock TFR catalogue build against the Carrick velocity field and the
    2MTF sky distribution to avoid recomputing the LOS velocities.
    """
    nsamples = len(RA_2MTF)

    # Convert Vext from ICRS to Galactic coordinates.
    Vext = np.asarray(Vext).reshape(1, 3)
    Vext_mag, Vext_RA, Vext_DEC = cartesian_to_radec(Vext).reshape(-1, )
    Vext_l, Vext_b = radec_to_galactic(Vext_RA, Vext_DEC)
    Vext_galactic = np.asanyarray([Vext_mag, Vext_l, Vext_b]).reshape(1, 3)
    Vext = radec_to_cartesian(Vext_galactic).reshape(-1, )

    truths = {"a": a_TF, "b": b_TF, "e_mu": sigma_TF, "sigma_v": sigma_v,
              "Vext": Vext,
              "mean_eta": mean_eta, "std_eta": std_eta,
              "mean_mag": mean_mag, "std_mag": std_mag,
              }

    gen = np.random.default_rng(seed)

    # The Carrick box is in the Galactic coordinates.
    l, b = radec_to_galactic(RA_2MTF, DEC_2MTF)
    gal_phi = np.deg2rad(l)
    gal_theta = np.pi / 2 - np.deg2rad(b)

    # Sample the linewidth of each galaxy from a Gaussian distribution to mimic
    # the MNR procedure.
    eta_true = gen.normal(mean_eta, std_eta, nsamples)
    eta_obs = gen.normal(eta_true, mean_e_eta)

    # Subtract the mean of the observed linewidths, so that they are
    # centered around zero. For consistency subtract from both observed
    # and true values.
    eta_mean_sampled = np.mean(eta_obs)
    eta_true -= eta_mean_sampled
    eta_obs -= eta_mean_sampled

    # Sample the magnitude from some Gaussian distribution to replicate MNR.
    mag_true = gen.normal(mean_mag, std_mag, nsamples)
    mag_obs = gen.normal(mag_true, mean_e_mag)

    # Calculate the 'true' distance modulus and redshift from the TFR distance.
    mu_TFR = mag_true - (a_TF + b_TF * eta_true)
    mu_true = gen.normal(mu_TFR, sigma_TF)

    # Convert the true distance modulus to true distance and cosmological
    # redshift.
    r = distmod2dist(mu_true, Om0)
    zcosmo = distmod2redshift(mu_true, Om0)

    # Calculate the Cartesian coordinates of each galaxy. This is initially
    # centered at (0, 0, 0).
    pos = r * np.asarray([
        np.sin(gal_theta) * np.cos(gal_phi),
        np.sin(gal_theta) * np.sin(gal_phi),
        np.cos(gal_theta)])
    pos = pos.T

    pos_box = pos / boxsize + 0.5
    vel = evaluate_cartesian_regular(
        velocity_field[0], velocity_field[1], velocity_field[2],
        pos=pos_box, smooth_scales=None, method="cubic")
    vel = beta * np.vstack(vel).T

    for i in range(3):
        vel[:, i] += Vext[i]

    # Compute the radial velocity.
    Vr = np.sum(vel * pos, axis=1) / np.linalg.norm(pos, axis=1)

    # The true redshift of the source.
    zCMB_true = (1 + zcosmo) * (1 + Vr / SPEED_OF_LIGHT) - 1
    zCMB_obs = gen.normal(zCMB_true, sigma_v / SPEED_OF_LIGHT)

    # These galaxies will be masked out when LOS is read it because they are
    # too far away.
    distance_mask = r < 125
    truths["distance_mask"] = distance_mask

    sample = {"RA": RA_2MTF,
              "DEC": DEC_2MTF,
              "z_CMB": zCMB_obs,
              "eta": eta_obs,
              "mag": mag_obs,
              "e_eta": np.ones(nsamples) * mean_e_eta,
              "e_mag": np.ones(nsamples) * mean_e_mag,
              "r": r,
              "distmod_true": mu_true,
              "distmod_TFR": mu_TFR}

    return sample, truths
