#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Created on Wed Apr 26 09:35:10 2023

@author: Christian Gimeno Amo (gimenoc@ifca.unican.es)

This module store all the necessary functions to do inpainting.

Dependencies:

- numpy
- healpy

"""

import numpy as np
import healpy as hp

def Z(Map, L, Mask_T, Pol = False, TPol = False, Mask_Pol = None, Fields = 3):
    
    """ This function computes the vector Z from the data and Cholesky decomposition.

    Parameters
    ----------
    Map : str or np.array
        Path to the map or map itself.
    L : np.array
        Block of the cholesky decomposition that contains the correlations between unmasked pixels.
    Mask_T : str or np.array
        Path to the mask or mask itself.
    Pol : bool, optional
        If QU maps are provided. The default is False.
    TPol : bool, optional
        If TQU maps are provided. The default is False.
    Mask_Pol : str or np.array, optional
        Path to the polarization mask or mask itself. The default is None.
    Fields : int, optional
        Number of fields. The default is 3.

    Returns
    -------
    z : np.array
        Vector Z.

    """

    if type(Map) == str:
        if TPol == False and Pol == False:
            Map = hp.read_map(Map, field = 0)
        else:
            F = [(0,1,2) if Pol and TPol else (0,1) if Pol and Fields == 3 else (1, 2)][0]
            Map = hp.read_map(Map, field = F)
    else:
        Map = Map

    if type(Mask_T) == str:
        Mask_T = hp.read_map(Mask_T)
    else:
        Mask_T = Mask_T

    Ind_T = np.where(Mask_T == 0)[0]
    
    if Pol == True and Mask_Pol is None:
        print("Error: Mask_Pol must be provided if Pol = True")
    elif Pol == True and type(Mask_Pol) == str:
        Mask_Pol = hp.read_map(Mask_Pol)
        Ind_Pol = np.where(Mask_Pol == 0)[0]
    elif Pol == True and isinstance(Mask_Pol, np.ndarray):
        Mask_Pol = Mask_Pol
        Ind_Pol = np.where(Mask_Pol == 0)[0]

    if Pol:
        if TPol:
            d1 = np.delete(Map[0], Ind_T, axis = 0)
            d2 = np.delete(Map[1], Ind_Pol, axis = 0)
            d3 = np.delete(Map[2], Ind_Pol, axis = 0)

            d = np.concatenate((d1, d2, d3), axis = 0)
            z = np.zeros(len(d))
            z[0] = d[0]/L[0, 0]

            for i in range(1, len(d)):
                vec = z[:i]*L[i, :i]
                z[i] = (d[i]-np.sum(vec))/L[i, i]
        else:
            d1 = np.delete(Map[0], Ind_Pol, axis = 0)
            d2 = np.delete(Map[1], Ind_Pol, axis = 0)

            d = np.concatenate((d1, d2), axis = 0)
            z = np.zeros(len(d))
            z[0] = d[0]/L[0, 0]

            for i in range(1, len(d)):
                vec = z[:i]*L[i, :i]
                z[i] = (d[i]-np.sum(vec))/L[i, i]
    else:
        d = np.delete(Map, Ind_T, axis = 0)

        z = np.zeros(len(d))
        z[0] = d[0]/L[0, 0]

        for i in range(1, len(d)):
            vec = z[:i]*L[i, :i]
            z[i] = (d[i]-np.sum(vec))/L[i, i]

    return z

def Inpaint(Map, R, L_bar, z, T_Mask, z_bar = None, Cons_Uncons = False, Pol = False, Pol_Mask = None, TPol = False, Fields = 3):

    """ This function inpaints the map given the Z vector, which contains the constrains of the unmasked pixels.

    Parameters
    ----------
    Map : str or np.array
        Path to the map or map itself.
    R : np.array
        Block of the cholesky decomposition that contains the correlations between masked and unmasked pixels.
    L_bar : np.array
        Block of the cholesky decomposition that contains the correlations between masked pixels.
    z : np.array
        Vector Z.
    T_Mask : str or np.array
        Path to the mask or mask itself.
    z_bar : np.array, optional
        Normally distributed vector of lenght equal to the number of masked pixels. The default is None, which means that is generated by the function itself.
    Cons_Uncons : bool, optional
        If True, the function returns the unconstrained map, the constrained map, and the inpainted map. The default is False.
    Pol : bool, optional
        If True QU maps are inpainted. The default is False.
    Pol_Mask : str or np.array, optional
        Path to the polarization mask or mask itself. The default is None.
    TPol : bool, optional
        If True TQU maps are inpainted. The default is False.
    Fields : int, optional
        Number of fields. The default is 3.

    Returns
    -------
    T_Cons_Part : np.array
        Constrained part of the T inpainted map. Only returned if Cons_Uncons = True. It is non stochastic.
    T_Uncons_Part : np.array
        Unconstrained part of the T inpainted map. Only returned if Cons_Uncons = True. It is stochastic.
    Inpainted_T_Map : np.array
        Inpainted T map.
    Q_Cons_Part : np.array
        Constrained part of the Q inpainted map. Only returned if Cons_Uncons = True. It is non stochastic.
    Q_Uncons_Part : np.array
        Unconstrained part of the Q inpainted map. Only returned if Cons_Uncons = True. It is stochastic.
    Inpainted_Q_Map : np.array
        Inpainted Q map.
    U_Cons_Part : np.array
        Constrained part of the U inpainted map. Only returned if Cons_Uncons = True. It is non stochastic.
    U_Uncons_Part : np.array
        Unconstrained part of the U inpainted map. Only returned if Cons_Uncons = True. It is stochastic.
    Inpainted_U_Map : np.array
        Inpainted U map.
    zbar : np.array
        Normally distributed vector of lenght equal to the number of masked pixels. If z_bar != None, output == input. 
        This scenario is useful if we want to inpaint a map given a certain seed.
    
    """

    if type(Map) == str:
        if Pol == False and TPol == False:
            Map = hp.read_map(Map, field = 0)
        else:
            F = [(0,1,2) if Pol and TPol else (0,1) if Pol and Fields == 3 else (1, 2)][0]
            Map = hp.read_map(Map, field = F)
    else:
        Map = Map

    if type(T_Mask) == str:
        Mask_T = hp.read_map(T_Mask)
    else:
        Mask_T = T_Mask

    Ind_T = np.where(Mask_T == 0)[0]
    
    if Pol == True and Pol_Mask == None:
        print("Error: Polarization Mask must be provided if Pol = True")
    elif Pol == True and type(Pol_Mask) == str:
        Mask_Pol = hp.read_map(Pol_Mask)
        Ind_Pol = np.where(Mask_Pol == 0)[0]
    elif Pol == True and isinstance(Pol_Mask, np.ndarray):
        Mask_Pol = Pol_Mask
        Ind_Pol = np.where(Mask_Pol == 0)[0]
        
    if z_bar is not None:
        zbar = np.load(z_bar)
    else:
        zbar = [np.random.normal(0, 1, size = len(Ind_T)+2*len(Ind_Pol)) if Pol == True and TPol == True else np.random.normal(0, 1, size = 2*len(Ind_Pol)) if TPol == False and Pol == True else np.random.normal(0, 1, size = len(Ind_T))][0]

    Constrained = np.dot(R, z)
    Unconstrained = np.dot(L_bar, zbar)
    dbar = Constrained+Unconstrained

    if Pol:
        if TPol:
            Inpainted_T_Map = np.copy(Map[0])
            Inpainted_Q_Map = np.copy(Map[1])
            Inpainted_U_Map = np.copy(Map[2])

            Inpainted_T_Map[Ind_T] = dbar[:len(Ind_T)]
            Inpainted_Q_Map[Ind_Pol] = dbar[len(Ind_T):(len(Ind_T)+len(Ind_Pol))]
            Inpainted_U_Map[Ind_Pol] = dbar[(len(Ind_T)+len(Ind_Pol)):]

            if Cons_Uncons:
                T_Cons_Part = np.copy(Map[0])
                T_Uncons_Part = np.zeros(len(Map[0]))
                Q_Cons_Part = np.copy(Map[1])
                Q_Uncons_Part = np.zeros(len(Map[1]))
                U_Cons_Part = np.copy(Map[2])
                U_Uncons_Part = np.zeros(len(Map[2]))
                T_Cons_Part[Ind_T] = Constrained[:len(Ind_T)]
                T_Uncons_Part[Ind_T] = Unconstrained[:len(Ind_T)]
                Q_Cons_Part[Ind_Pol] = Constrained[len(Ind_T):(len(Ind_T)+len(Ind_Pol))]
                Q_Uncons_Part[Ind_Pol] = Unconstrained[len(Ind_T):(len(Ind_T)+len(Ind_Pol))]
                U_Cons_Part[Ind_Pol] = Constrained[(len(Ind_T)+len(Ind_Pol)):]
                U_Uncons_Part[Ind_Pol] = Unconstrained[(len(Ind_T)+len(Ind_Pol)):]

                return (T_Cons_Part, T_Uncons_Part, Inpainted_T_Map, Q_Cons_Part, Q_Uncons_Part, Inpainted_Q_Map, U_Cons_Part, U_Uncons_Part, Inpainted_U_Map, zbar)
            else:
                return (Inpainted_T_Map, Inpainted_Q_Map, Inpainted_U_Map, zbar)
        else:
            Inpainted_Q_Map = np.copy(Map[0])
            Inpainted_U_Map = np.copy(Map[1])
            Inpainted_Q_Map[Ind_Pol] = dbar[:len(Ind_Pol)]
            Inpainted_U_Map[Ind_Pol] = dbar[len(Ind_Pol):]

            if Cons_Uncons:
                Q_Cons_Part = np.copy(Map[0])
                Q_Uncons_Part = np.zeros(len(Map[0]))
                U_Cons_Part = np.copy(Map[1])
                U_Uncons_Part = np.zeros(len(Map[1]))
                Q_Cons_Part[Ind_Pol] = Constrained[:len(Ind_Pol)]
                Q_Uncons_Part[Ind_Pol] = Unconstrained[:len(Ind_Pol)]
                U_Cons_Part[Ind_Pol] = Constrained[len(Ind_Pol):]
                U_Uncons_Part[Ind_Pol] = Unconstrained[len(Ind_Pol):]

                return (Q_Cons_Part, Q_Uncons_Part, Inpainted_Q_Map, U_Cons_Part, U_Uncons_Part, Inpainted_U_Map, zbar)
            else:
                return (Inpainted_Q_Map, Inpainted_U_Map, zbar)

    else:
        Inpainted_T_Map = np.copy(Map)
        Inpainted_T_Map[Ind_T] = dbar

        if Cons_Uncons:
            T_Cons_Part = np.copy(Map)
            T_Uncons_Part = np.zeros(len(Map))
            T_Cons_Part[Ind_T] = Constrained
            T_Uncons_Part[Ind_T] = Unconstrained

            return (T_Cons_Part, T_Uncons_Part, Inpainted_T_Map, zbar)
        else:
            return (Inpainted_T_Map, zbar)

__docformat__ = "numpy"
