import os
import time
import uuid
import yaml
import logging
import shutil
import numpy as np
import pandas as pd
import multiprocessing as mp
from functools import partial
from astropy.time import Time

from .config import Config
from .config import Configuration
from .clusters import find_clusters, filter_clusters_by_length
from .cell import Cell
from .orbit import TestOrbit
from .orbits import Orbits
from .orbits import generateEphemeris
from .orbits import initialOrbitDetermination
from .orbits import differentialCorrection
from .orbits import mergeAndExtendOrbits
from .observatories import getObserverState
from .utils import _initWorker
from .utils import _checkParallel

os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

logger = logging.getLogger("thor")

__all__ = [
    "rangeAndShift_worker",
    "rangeAndShift",
    "clusterVelocity",
    "clusterVelocity_worker",
    "clusterAndLink",
    "runTHOROrbit",
    "runTHOR",
]

def rangeAndShift_worker(observations, ephemeris, cell_area=10):

    assert len(observations["mjd_utc"].unique()) == 1
    assert len(ephemeris["mjd_utc"].unique()) == 1
    assert observations["mjd_utc"].unique()[0] == ephemeris["mjd_utc"].unique()[0]
    observation_time = observations["mjd_utc"].unique()[0]

    # Create Cell centered on the sky-plane location of the
    # test orbit
    cell = Cell(
        ephemeris[["RA_deg", "Dec_deg"]].values[0],
        observation_time,
        area=cell_area,
    )

    # Grab observations within cell
    cell.getObservations(observations)

    if len(cell.observations) != 0:

        # Create test orbit with state of orbit at visit time
        test_orbit = TestOrbit(
            ephemeris[["obj_x", "obj_y", "obj_z", "obj_vx", "obj_vy", "obj_vz"]].values[0],
            observation_time
        )

        # Prepare rotation matrices
        test_orbit.prepare()

        # Apply rotation matrices and transform observations into the orbit's
        # frame of motion.
        test_orbit.applyToObservations(cell.observations)

        projected_observations = cell.observations

    else:

        projected_observations = pd.DataFrame()

    return projected_observations

def clusterVelocity(
        obs_ids,
        x,
        y,
        dt,
        vx,
        vy,
        eps=0.005,
        min_obs=5,
        min_arc_length=1.0,
        alg="hotspot_2d",
    ):
    """
    Clusters THOR projection with different velocities
    in the projection plane using `~scipy.cluster.DBSCAN`.
    Parameters
    ----------
    obs_ids : `~numpy.ndarray' (N)
        Observation IDs.
    x : `~numpy.ndarray' (N)
        Projection space x coordinate in degrees or radians.
    y : `~numpy.ndarray' (N)
        Projection space y coordinate in degrees or radians.
    dt : `~numpy.ndarray' (N)
        Change in time from 0th exposure in units of MJD.
    vx : `~numpy.ndarray' (N)
        Projection space x velocity in units of degrees or radians per day in MJD.
    vy : `~numpy.ndarray' (N)
        Projection space y velocity in units of degrees or radians per day in MJD.
    eps : float, optional
        The maximum distance between two samples for them to be considered
        as in the same neighborhood.
        See: http://scikit-learn.org/stable/modules/generated/sklearn.cluster.dbscan.html
        [Default = 0.005]
    min_obs : int, optional
        The number of samples (or total weight) in a neighborhood for a
        point to be considered as a core point. This includes the point itself.
        See: http://scikit-learn.org/stable/modules/generated/sklearn.cluster.dbscan.html
        [Default = 5]
    min_arc_length : float, optional
        Minimum arc length in units of days for a cluster to be accepted.

    Returns
    -------
    list
        If clusters are found, will return a list of numpy arrays containing the
        observation IDs for each cluster. If no clusters are found, will return np.NaN.
    """
    logger.debug(f"cluster: vx={vx} vy={vy} n_obs={len(obs_ids)}")
    xx = x - vx * dt
    yy = y - vy * dt

    X = np.stack((xx, yy), 1)

    clusters = find_clusters(X, eps, min_obs, alg=alg)
    clusters = filter_clusters_by_length(
        clusters, dt, min_obs, min_arc_length,
    )

    cluster_ids = []
    for cluster in clusters:
        cluster_ids.append(obs_ids[cluster])

    if len(cluster_ids) == 0:
        cluster_ids = np.NaN

    return cluster_ids

def clusterVelocity_worker(
        vx,
        vy,
        obs_ids=None,
        x=None,
        y=None,
        dt=None,
        eps=None,
        min_obs=None,
        min_arc_length=None,
        alg=None
    ):
    """
    Helper function to multiprocess clustering.

    """
    cluster_ids = clusterVelocity(
        obs_ids,
        x,
        y,
        dt,
        vx,
        vy,
        eps=eps,
        min_obs=min_obs,
        min_arc_length=min_arc_length,
        alg=alg
    )
    return cluster_ids

def rangeAndShift(
        observations,
        orbit,
        cell_area=10,
        backend="PYOORB",
        backend_kwargs={},
        num_jobs=1,
        parallel_backend="mp"
    ):
    """
    Propagate the orbit to all observation times in observations. At each epoch gather a circular region of observations of size cell_area
    centered about the location of the orbit on the sky-plane. Transform and project each of the gathered observations into
    the frame of motion of the test orbit.

    Parameters
    ----------
    observations : `~pandas.DataFrame`
        DataFrame containing preprocessed observations.
            Should contain the following columns:
                obs_id : observation IDs
                RA_deg : Right Ascension in degrees.
                Dec_deg : Declination in degrees.
                RA_sigma_deg : 1-sigma uncertainty for Right Ascension in degrees.
                Dec_sigma_deg : 1-sigma uncertainty for Declination in degrees.
                observatory_code : MPC observatory code
    orbit : `~numpy.ndarray` (6)
        Orbit to propagate. If backend is 'THOR', then these orbits must be expressed
        as heliocentric ecliptic cartesian elements. If backend is 'PYOORB' orbits may be
        expressed in keplerian, cometary or cartesian elements.
    cell_area : float, optional
        Cell's area in units of square degrees.
        [Default = 10]
    backend : {'THOR', 'PYOORB'}, optional
        Which backend to use.
    backend_kwargs : dict, optional
        Settings and additional parameters to pass to selected
        backend.
    num_jobs : int, optional
        Number of jobs to launch.
    parallel_backend : str, optional
        Which parallelization backend to use {'ray', 'mp'}. Defaults to using Python's multiprocessing
        module ('mp').

    Returns
    -------
    projected_observations : {`~pandas.DataFrame`, -1}
        Observations dataframe (from cell.observations) with columns containing
        projected coordinates.
    """
    time_start = time.time()
    logger.info("Running range and shift...")
    logger.info("Assuming r = {} au".format(orbit.cartesian[0, :3]))
    logger.info("Assuming v = {} au per day".format(orbit.cartesian[0, 3:]))

    # Build observers dictionary: keys are observatory codes with exposure times (as astropy.time objects)
    # as values
    observers = {}
    for code in observations["observatory_code"].unique():
        observers[code] = Time(
            observations[observations["observatory_code"].isin([code])]["mjd_utc"].unique(),
            format="mjd",
            scale="utc"
        )

    # Propagate test orbit to all times in observations
    ephemeris = generateEphemeris(
        orbit,
        observers,
        backend=backend,
        backend_kwargs=backend_kwargs,
        chunk_size=1,
        num_jobs=1,
        parallel_backend=parallel_backend
    )
    if backend == "FINDORB":

        observer_states = []
        for observatory_code, observation_times in observers.items():
            observer_states.append(
                getObserverState(
                    [observatory_code],
                    observation_times,
                    frame='ecliptic',
                    origin='heliocenter',
                )
            )

        observer_states = pd.concat(observer_states)
        observer_states.reset_index(
            inplace=True,
            drop=True
        )
        ephemeris = ephemeris.join(observer_states[["obs_x", "obs_y", "obs_z", "obs_vx", "obs_vy", "obs_vz"]])

    velocity_cols = []
    if backend != "PYOORB":
        velocity_cols = ["obs_vx", "obs_vy", "obs_vz"]

    observations = observations.merge(
        ephemeris[["mjd_utc", "observatory_code", "obs_x", "obs_y", "obs_z"] + velocity_cols],
        left_on=["mjd_utc", "observatory_code"],
        right_on=["mjd_utc", "observatory_code"]
    )

    # Split the observations into a single dataframe per unique observatory code and observation time
    # Basically split the observations into groups of unique exposures
    observations_grouped = observations.groupby(by=["observatory_code", "mjd_utc"])
    observations_split = [observations_grouped.get_group(g) for g in observations_grouped.groups]

    # Do the same for the test orbit's ephemerides
    ephemeris_grouped = ephemeris.groupby(by=["observatory_code", "mjd_utc"])
    ephemeris_split = [ephemeris_grouped.get_group(g) for g in ephemeris_grouped.groups]

    parallel, num_workers = _checkParallel(num_jobs, parallel_backend)
    if parallel:
        if parallel_backend == "ray":
            import ray
            if not ray.is_initialized():
                ray.init(address="auto")

            rangeAndShift_worker_ray = ray.remote(rangeAndShift_worker)
            rangeAndShift_worker_ray = rangeAndShift_worker_ray.options(
                num_returns=1,
                num_cpus=1
            )

            p = []
            for observations_i, ephemeris_i in zip(observations_split, ephemeris_split):
                p.append(
                    rangeAndShift_worker_ray.remote(
                        observations_i,
                        ephemeris_i,
                        cell_area=cell_area
                    )
                )
            projected_dfs = ray.get(p)

        else: # parallel_backend == "mp"
            p = mp.Pool(
                processes=num_workers,
                initializer=_initWorker,
            )
            projected_dfs = p.starmap(
                partial(
                    rangeAndShift_worker,
                    cell_area=cell_area
                ),
                zip(
                    observations_split,
                    ephemeris_split,
                )
            )
            p.close()

    else:
        projected_dfs = []
        for observations_i, ephemeris_i in zip(observations_split, ephemeris_split):
            projected_df = rangeAndShift_worker(
                observations_i,
                ephemeris_i,
                cell_area=cell_area
            )
            projected_dfs.append(projected_df)

    projected_observations = pd.concat(projected_dfs)
    if len(projected_observations) > 0:
        projected_observations.sort_values(by=["mjd_utc", "observatory_code"], inplace=True)
        projected_observations.reset_index(inplace=True, drop=True)
    else:
        projected_observations = pd.DataFrame(
            columns=[
                'obs_id', 'mjd_utc', 'RA_deg', 'Dec_deg', 'RA_sigma_deg',
                'Dec_sigma_deg', 'observatory_code', 'obs_x', 'obs_y', 'obs_z', 'obj_x',
                'obj_y', 'obj_z', 'theta_x_deg', 'theta_y_deg'
            ]
        )

    time_end = time.time()
    logger.info("Found {} observations.".format(len(projected_observations)))
    logger.info("Range and shift completed in {:.3f} seconds.".format(time_end - time_start))

    return projected_observations

def clusterAndLink(
        observations,
        vx_range=[-0.1, 0.1],
        vy_range=[-0.1, 0.1],
        vx_bins=100,
        vy_bins=100,
        vx_values=None,
        vy_values=None,
        eps=0.005,
        min_obs=5,
        min_arc_length=1.0,
        alg="dbscan",
        num_jobs=1,
        parallel_backend="mp"
    ):
    """
    Cluster and link correctly projected (after ranging and shifting)
    detections.

    Parameters
    ----------
    observations : `~pandas.DataFrame`
        DataFrame containing post-range and shift observations.
    vx_range : {None, list or `~numpy.ndarray` (2)}
        Maximum and minimum velocity range in x.
        Will not be used if vx_values are specified.
        [Default = [-0.1, 0.1]]
    vy_range : {None, list or `~numpy.ndarray` (2)}
        Maximum and minimum velocity range in y.
        Will not be used if vy_values are specified.
        [Default = [-0.1, 0.1]]
    vx_bins : int, optional
        Length of x-velocity grid between vx_range[0]
        and vx_range[-1]. Will not be used if vx_values are
        specified.
        [Default = 100]
    vy_bins: int, optional
        Length of y-velocity grid between vy_range[0]
        and vy_range[-1]. Will not be used if vy_values are
        specified.
        [Default = 100]
    vx_values : {None, `~numpy.ndarray`}, optional
        Values of velocities in x at which to cluster
        and link.
        [Default = None]
    vy_values : {None, `~numpy.ndarray`}, optional
        Values of velocities in y at which to cluster
        and link.
        [Default = None]
    eps : float, optional
        The maximum distance between two samples for them to be considered
        as in the same neighborhood.
        See: http://scikit-learn.org/stable/modules/generated/sklearn.cluster.dbscan.html
        [Default = 0.005]
    min_obs : int, optional
        The number of samples (or total weight) in a neighborhood for a
        point to be considered as a core point. This includes the point itself.
        See: http://scikit-learn.org/stable/modules/generated/sklearn.cluster.dbscan.html
        [Default = 5]
    alg: str
        Algorithm to use. Can be "dbscan" or "hotspot_2d".
    num_jobs : int, optional
        Number of jobs to launch.
    parallel_backend : str, optional
        Which parallelization backend to use {'ray', 'mp'}. Defaults to using Python's multiprocessing
        module ('mp').

    Returns
    -------
    clusters : `~pandas.DataFrame`
        DataFrame with the cluster ID, the number of observations, and the x and y velocity.
    cluster_members : `~pandas.DataFrame`
        DataFrame containing the cluster ID and the observation IDs of its members.

    Notes
    -----
    The algorithm chosen can have a big impact on performance and accuracy.

    alg="dbscan" uses the DBSCAN algorithm of Ester et. al. It's relatively slow
    but works with high accuracy; it is certain to find all clusters with at
    least min_obs points that are separated by at most eps.

    alg="hotspot_2d" is much faster (perhaps 10-20x faster) than dbscan, but it
    may miss some clusters, particularly when points are spaced a distance of 'eps'
    apart.
    """
    time_start_cluster = time.time()
    logger.info("Running velocity space clustering...")

    # Extract useful quantities
    obs_ids = observations["obs_id"].values
    theta_x = observations["theta_x_deg"].values
    theta_y = observations["theta_y_deg"].values
    mjd = observations["mjd_utc"].values

    # Select detections in first exposure
    first = np.where(mjd == mjd.min())[0]
    mjd0 = mjd[first][0]
    dt = mjd - mjd0

    if vx_values is None and vx_range is not None:
        vx = np.linspace(*vx_range, num=vx_bins)
    elif vx_values is None and vx_range is None:
        raise ValueError("Both vx_values and vx_range cannot be None.")
    else:
        vx = vx_values
        vx_range = [vx_values[0], vx_values[-1]]
        vx_bins = len(vx)

    if vy_values is None and vy_range is not None:
        vy = np.linspace(*vy_range, num=vy_bins)
    elif vy_values is None and vy_range is None:
        raise ValueError("Both vy_values and vy_range cannot be None.")
    else:
        vy = vy_values
        vy_range = [vy_values[0], vy_values[-1]]
        vy_bins = len(vy)

    if vx_values is None and vy_values is None:
        vxx, vyy = np.meshgrid(vx, vy)
        vxx = vxx.flatten()
        vyy = vyy.flatten()
    elif vx_values is not None and vy_values is not None:
        vxx = vx
        vyy = vy
    else:
        raise ValueError("")

    logger.debug("X velocity range: {}".format(vx_range))
    if vx_values is not None:
        logger.debug("X velocity values: {}".format(vx_bins))
    else:
        logger.debug("X velocity bins: {}".format(vx_bins))

    logger.debug("Y velocity range: {}".format(vy_range))
    if vy_values is not None:
        logger.debug("Y velocity values: {}".format(vy_bins))
    else:
        logger.debug("Y velocity bins: {}".format(vy_bins))
    if vx_values is not None:
        logger.debug("User defined x velocity values: True")
    else:
        logger.debug("User defined x velocity values: False")
    if vy_values is not None:
        logger.debug("User defined y velocity values: True")
    else:
        logger.debug("User defined y velocity values: False")

    if vx_values is None and vy_values is None:
        logger.debug("Velocity grid size: {}".format(vx_bins * vy_bins))
    else:
        logger.debug("Velocity grid size: {}".format(vx_bins))
    logger.info("Max sample distance: {}".format(eps))
    logger.info("Minimum samples: {}".format(min_obs))

    possible_clusters = []
    parallel, num_workers = _checkParallel(num_jobs, parallel_backend)
    if parallel:
        if parallel_backend == "ray":
            import ray
            if not ray.is_initialized():
                ray.init(address="auto")

            clusterVelocity_worker_ray = ray.remote(clusterVelocity_worker)
            clusterVelocity_worker_ray = clusterVelocity_worker_ray.options(
                num_returns=1,
                num_cpus=1
            )

            # Put all arrays (which can be large) in ray's
            # local object store ahead of time
            obs_ids_oid = ray.put(obs_ids)
            theta_x_oid = ray.put(theta_x)
            theta_y_oid = ray.put(theta_y)
            dt_oid = ray.put(dt)

            p = []
            for vxi, vyi in zip(vxx, vyy):
                p.append(
                    clusterVelocity_worker_ray.remote(
                        vxi,
                        vyi,
                        obs_ids=obs_ids_oid,
                        x=theta_x_oid,
                        y=theta_y_oid,
                        dt=dt_oid,
                        eps=eps,
                        min_obs=min_obs,
                        min_arc_length=min_arc_length,
                        alg=alg
                    )
                )
            possible_clusters = ray.get(p)

        else: # parallel_backend == "mp"

            p = mp.Pool(
                processes=num_workers,
                initializer=_initWorker
            )
            possible_clusters = p.starmap(
                partial(
                    clusterVelocity_worker,
                    obs_ids=obs_ids,
                    x=theta_x,
                    y=theta_y,
                    dt=dt,
                    eps=eps,
                    min_obs=min_obs,
                    min_arc_length=min_arc_length,
                    alg=alg
                ),
                zip(vxx, vyy)
            )
            p.close()

    else:
        possible_clusters = []
        for vxi, vyi in zip(vxx, vyy):
            possible_clusters.append(
                clusterVelocity(
                    obs_ids,
                    theta_x,
                    theta_y,
                    dt,
                    vxi,
                    vyi,
                    eps=eps,
                    min_obs=min_obs,
                    min_arc_length=min_arc_length,
                    alg=alg
                )
            )
    time_end_cluster = time.time()
    logger.info("Clustering completed in {:.3f} seconds.".format(time_end_cluster - time_start_cluster))

    logger.info("Restructuring clusters...")
    time_start_restr = time.time()

    possible_clusters = pd.DataFrame({"clusters": possible_clusters})

    # Remove empty clusters
    possible_clusters = possible_clusters[~possible_clusters["clusters"].isna()]

    if len(possible_clusters) != 0:
        ### The following code is a little messy, its a lot of pandas dataframe manipulation.
        ### I have tried doing an overhaul wherein the clusters and cluster_members dataframe are created per
        ### velocity combination in the clusterVelocity function. However, this adds an overhead in that function
        ### of ~ 1ms. So clustering 90,000 velocities takes 90 seconds longer which on small datasets is problematic.
        ### On large datasets, the effect is not as pronounced because the below code takes a while to run due to
        ### in-memory pandas dataframe restructuring.

        # Make DataFrame with cluster velocities so we can figure out which
        # velocities yielded clusters, add names to index so we can enable the join
        cluster_velocities = pd.DataFrame({"vtheta_x": vxx, "vtheta_y": vyy})
        cluster_velocities.index.set_names("velocity_id", inplace=True)

        # Split lists of cluster ids into one column per cluster for each different velocity
        # then stack the result
        possible_clusters = pd.DataFrame(
            possible_clusters["clusters"].values.tolist(),
            index=possible_clusters.index
        )
        possible_clusters = pd.DataFrame(possible_clusters.stack())
        possible_clusters.rename(
            columns={0: "obs_ids"},
            inplace=True
        )
        possible_clusters = pd.DataFrame(possible_clusters["obs_ids"].values.tolist(), index=possible_clusters.index)

        # Drop duplicate clusters
        possible_clusters.drop_duplicates(inplace=True)

        # Set index names
        possible_clusters.index.set_names(["velocity_id", "cluster_id"], inplace=True)

        # Reset index
        possible_clusters.reset_index(
            "cluster_id",
            drop=True,
            inplace=True
        )
        possible_clusters["cluster_id"] = [str(uuid.uuid4().hex) for i in range(len(possible_clusters))]

        # Make clusters DataFrame
        clusters = possible_clusters.join(cluster_velocities)
        clusters.reset_index(drop=True, inplace=True)
        clusters = clusters[["cluster_id", "vtheta_x", "vtheta_y"]]

        # Make cluster_members DataFrame
        cluster_members = possible_clusters.reset_index(drop=True).copy()
        cluster_members.index = cluster_members["cluster_id"]
        cluster_members.drop("cluster_id", axis=1, inplace=True)
        cluster_members = pd.DataFrame(cluster_members.stack())
        cluster_members.rename(columns={0: "obs_id"}, inplace=True)
        cluster_members.reset_index(inplace=True)
        cluster_members.drop("level_1", axis=1, inplace=True)

        # Calculate arc length and add it to the clusters dataframe
        cluster_members_time = cluster_members.merge(
            observations[["obs_id", "mjd_utc"]],
            on="obs_id",
            how="left"
        )
        clusters_time = cluster_members_time.groupby(
            by=["cluster_id"])["mjd_utc"].apply(lambda x: x.max() - x.min()).to_frame()
        clusters_time.reset_index(
            inplace=True
        )
        clusters_time.rename(
            columns={"mjd_utc" : "arc_length"},
            inplace=True
        )
        clusters = clusters.merge(
            clusters_time[["cluster_id", "arc_length"]],
            on="cluster_id",
            how="left",
        )

    else:
        cluster_members = pd.DataFrame(columns=["cluster_id", "obs_id"])
        clusters = pd.DataFrame(columns=["cluster_id", "vtheta_x", "vtheta_y", "arc_length"])

    time_end_restr = time.time()
    logger.info("Restructuring completed in {:.3f} seconds.".format(time_end_restr - time_start_restr))
    logger.info("Found {} clusters.".format(len(clusters)))
    logger.info("Clustering and restructuring completed in {:.3f} seconds.".format(time_end_restr - time_start_cluster))

    return clusters, cluster_members

def runTHOROrbit(
        preprocessed_observations,
        orbit,
        range_shift_config=Config.RANGE_SHIFT_CONFIG,
        cluster_link_config=Config.CLUSTER_LINK_CONFIG,
        iod_config=Config.IOD_CONFIG,
        od_config=Config.OD_CONFIG,
        odp_config=Config.ODP_CONFIG,
        out_dir=None,
        if_exists="continue",
        logging_level=logging.INFO
    ):
    logger = logging.getLogger("thor")
    logger.setLevel(logging_level)

    # Build the configuration class which stores the run parameters
    config = Configuration(
        range_shift_config=range_shift_config,
        cluster_link_config=cluster_link_config,
        iod_config=iod_config,
        od_config=od_config,
        odp_config=odp_config
    )
    status = {
        "rangeAndShift" : False,
        "clusterAndLink" : False,
        "initialOrbitDetermination" : False,
        "differentialCorrection" : False,
        "mergeAndExtendOrbits" : False,
        "complete" : False
    }

    continue_ = False
    if out_dir is not None:
        if not os.path.exists(out_dir):
            os.mkdir(out_dir)
            logger.debug("Created {} directory.".format(out_dir))

        else:
            if if_exists == "continue":
                logger.warning("{} directory already exists, attempting to continue previous run.".format(out_dir))
                continue_ = True
            elif if_exists == "erase":
                logger.warning("{} directory already exists, removing previous results.".format(out_dir))
                shutil.rmtree(out_dir)
                os.mkdir(out_dir)
                logger.debug("Created {} directory.".format(out_dir))
            else:
                err = (
                    "if_exists should be one of {'continue', 'erase'}."
                )
                raise ValueError(err)

        file_handler = logging.FileHandler(
            os.path.join(out_dir, "thor.log"),
            encoding="utf-8",
            delay=False
        )
        file_handler.setLevel(logging.DEBUG)
        file_format = logging.Formatter(
            '%(asctime)s.%(msecs)03d [%(levelname)s] [%(thread)s] %(message)s (%(filename)s, %(funcName)s, %(lineno)d)',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        file_handler.setFormatter(file_format)
        logger.addHandler(file_handler)

        # The primary files which will be used to determine if the run
        # can be continued from a previous state and, if so, from where
        # to continue the run
        config_file = os.path.join(out_dir, "config.yml")
        test_orbit_file = os.path.join(out_dir, "test_orbit.csv")
        status_file = os.path.join(out_dir, "status.yml")
        config_eq = False
        test_orbit_eq = False
        save_orbit = True
        save_config = True

        if continue_:

            if not os.path.exists(config_file):
                logger.warning("No previous configuration file found.")
                save_config = True
            else:
                logger.info("Previous configuration file found. Comparing settings...")

                config_prev = Configuration.fromYaml(config_file)
                if config_prev != config:
                    logger.warning("Previous configuration does not match current configuration. Processing will not continue from previous state.")
                else:
                    config_eq = True
                    save_config = False
                    logger.info("Previous configuration matches current configuration.")

            if not os.path.exists(test_orbit_file):
                logger.warning("No previous test orbit file found.")
                save_orbit = True
            else:
                logger.info("Previous test orbit file found.")

                test_orbit_prev = Orbits.from_csv(
                    test_orbit_file,
                )
                if test_orbit_prev != orbit:
                    logger.warning("Previous test orbit does not match current test orbit.")
                else:
                    test_orbit_eq = True
                    save_orbit = False
                    logger.info("Previous test orbit matches current test orbit.")

            if not os.path.exists(status_file):
                logger.warning("No previous status file found.")
            else:
                if test_orbit_eq and config_eq:
                    with open(status_file, "r") as status_in:
                        status = yaml.load(status_in, Loader=yaml.FullLoader)
                    logger.info("Previous status file found.")

        if save_config:
            config.toYaml(config_file)
            logger.debug("Saved config.yml.")

        if save_orbit:
            orbit.to_csv(
                test_orbit_file
            )
            logger.debug("Saved test_orbit.csv.")

            if status["complete"]:
                logger.info("Orbit has already finished processing.")

    if not status["complete"]:
        if not status["rangeAndShift"]:
            projected_observations = rangeAndShift(
                preprocessed_observations,
                orbit,
                **range_shift_config
            )
            if out_dir is not None:
                projected_observations.to_csv(
                    os.path.join(out_dir, "projected_observations.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved projected_observations.csv.")

        else:
            logger.info("Range and shift completed previously.")
            projected_observations = pd.read_csv(
                os.path.join(out_dir, "projected_observations.csv"),
                index_col=False,
                dtype={"obs_id" : str},
                float_precision="round_trip"
            )
            logger.debug("Read projected_observations.csv.")

        status["rangeAndShift"] = True
        if out_dir is not None:
            with open(status_file, "w") as status_out:
                yaml.safe_dump(status, status_out)
            logger.debug("Updated status.yml.")

        if not status["clusterAndLink"]:
            clusters, cluster_members = clusterAndLink(
                projected_observations,
                **cluster_link_config
            )
            if out_dir is not None:
                clusters.to_csv(
                    os.path.join(out_dir, "clusters.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved clusters.csv.")

                cluster_members.to_csv(
                    os.path.join(out_dir, "cluster_members.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved cluster_members.csv.")
        else:
            logger.info("Clustering completed previously.")
            clusters = pd.read_csv(
                os.path.join(out_dir, "clusters.csv"),
                index_col=False
            )
            logger.debug("Read clusters.csv.")

            cluster_members = pd.read_csv(
                os.path.join(out_dir, "cluster_members.csv"),
                index_col=False,
                dtype={"obs_id" : str},
                float_precision="round_trip"
            )
            logger.debug("Read cluster_members.csv.")

        status["clusterAndLink"] = True
        if out_dir is not None:
            with open(status_file, "w") as status_out:
                yaml.safe_dump(status, status_out)
            logger.debug("Updated status.yml.")

        if not status["initialOrbitDetermination"]:
            iod_orbits, iod_orbit_members = initialOrbitDetermination(
                projected_observations,
                cluster_members,
                **iod_config
            )
            if out_dir is not None:
                Orbits.from_df(iod_orbits).to_csv(
                    os.path.join(out_dir, "iod_orbits.csv")
                )
                logger.debug("Saved iod_orbits.csv.")

                iod_orbit_members.to_csv(
                    os.path.join(out_dir, "iod_orbit_members.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved iod_orbit_members.csv.")
        else:
            logger.info("Initial orbit determination completed previously.")
            iod_orbits = Orbits.from_csv(
                os.path.join(out_dir, "iod_orbits.csv"),
            ).to_df(include_units=False)
            logger.debug("Read iod_orbits.csv.")

            iod_orbit_members = pd.read_csv(
                os.path.join(out_dir, "iod_orbit_members.csv"),
                index_col=False,
                dtype={"obs_id" : str},
                float_precision="round_trip"
            )
            logger.debug("Read iod_orbit_members.csv.")

        status["initialOrbitDetermination"] = True
        if out_dir is not None:
            with open(status_file, "w") as status_out:
                yaml.safe_dump(status, status_out)
            logger.debug("Updated status.yml.")

        iod_orbits = iod_orbits[["orbit_id", "epoch", "x", "y", "z", "vx", "vy", "vz"]]
        iod_orbit_members = iod_orbit_members[iod_orbit_members["outlier"] == 0][["orbit_id", "obs_id"]]
        iod_orbits = iod_orbits[iod_orbits["orbit_id"].isin(iod_orbit_members["orbit_id"].unique())]
        for df in [iod_orbits, iod_orbit_members]:
            df.reset_index(
                inplace=True,
                drop=True
            )

        if not status["differentialCorrection"]:
            od_orbits, od_orbit_members = differentialCorrection(
                iod_orbits,
                iod_orbit_members,
                projected_observations,
                **od_config
            )
            if out_dir is not None:
                Orbits.from_df(od_orbits).to_csv(
                    os.path.join(out_dir, "od_orbits.csv")
                )
                logger.debug("Saved od_orbits.csv.")

                od_orbit_members.to_csv(
                    os.path.join(out_dir, "od_orbit_members.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved od_orbit_members.csv.")
        else:
            logger.info("Differential correction completed previously.")
            od_orbits = Orbits.from_csv(
                os.path.join(out_dir, "od_orbits.csv"),
            ).to_df(include_units=False)
            logger.debug("Read od_orbits.csv.")

            od_orbit_members = pd.read_csv(
                os.path.join(out_dir, "od_orbit_members.csv"),
                index_col=False,
                dtype={"obs_id" : str},
                float_precision="round_trip"
            )
            logger.debug("Read od_orbit_members.csv.")

        status["differentialCorrection"] = True
        if out_dir is not None:
            with open(status_file, "w") as status_out:
                yaml.safe_dump(status, status_out)
            logger.debug("Updated status.yml.")

        od_orbit_members = od_orbit_members[od_orbit_members["outlier"] == 0][["orbit_id", "obs_id"]]
        od_orbits = od_orbits[od_orbits["orbit_id"].isin(od_orbit_members["orbit_id"].unique())]
        for df in [od_orbits, od_orbit_members]:
            df.reset_index(
                inplace=True,
                drop=True
            )

        if not status["mergeAndExtendOrbits"]:

            recovered_orbits, recovered_orbit_members = mergeAndExtendOrbits(
                od_orbits,
                od_orbit_members,
                projected_observations,
                **odp_config
            )
            if out_dir is not None:
                Orbits.from_df(recovered_orbits).to_csv(
                    os.path.join(out_dir, "recovered_orbits.csv")
                )
                logger.debug("Saved recovered_orbits.csv.")

                recovered_orbit_members.to_csv(
                    os.path.join(out_dir, "recovered_orbit_members.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved recovered_orbit_members.csv.")
        else:
            logger.info("Orbit extension and merging completed previously.")
            recovered_orbits = Orbits.from_csv(
                os.path.join(out_dir, "recovered_orbits.csv"),
            ).to_df(include_units=False)
            logger.debug("Read recovered_orbits.csv.")

            recovered_orbit_members = pd.read_csv(
                os.path.join(out_dir, "recovered_orbit_members.csv"),
                index_col=False,
                dtype={"obs_id" : str},
                float_precision="round_trip"
            )
            logger.debug("Read recovered_orbit_members.csv.")

        status["mergeAndExtendOrbits"] = True
        status["complete"] = True
        if out_dir is not None:
            with open(status_file, "w") as status_out:
                yaml.safe_dump(status, status_out)
            logger.debug("Updated status.yml.")

    else:
        logger.info("Orbit previously completed processing.")
        recovered_orbits = Orbits.from_csv(
            os.path.join(out_dir, "recovered_orbits.csv"),
        ).to_df(include_units=False)
        logger.debug("Read recovered_orbits.csv.")

        recovered_orbit_members = pd.read_csv(
            os.path.join(out_dir, "recovered_orbit_members.csv"),
            index_col=False,
            dtype={"obs_id" : str},
            float_precision="round_trip"
        )
        logger.debug("Read recovered_orbit_members.csv.")

    logger.removeHandler(file_handler)
    return recovered_orbits, recovered_orbit_members

def runTHOR(
        preprocessed_observations,
        test_orbits,
        range_shift_config=Config.RANGE_SHIFT_CONFIG,
        cluster_link_config=Config.CLUSTER_LINK_CONFIG,
        iod_config=Config.IOD_CONFIG,
        od_config=Config.OD_CONFIG,
        odp_config=Config.ODP_CONFIG,
        out_dir=None,
        if_exists="continue",
        logging_level=logger.info
    ):
    logger.setLevel(logging_level)

    # Connect to ray cluster if enabled
    enable_ray = False
    configs = [
        range_shift_config,
        cluster_link_config,
        iod_config,
        od_config,
        odp_config
    ]
    for conf in configs:
        if conf["parallel_backend"] == "ray":
            enable_ray = True

    if enable_ray:
        import ray
        if not ray.is_initialized():
            ray.init(address="auto")

    # Build the configuration class which stores the run parameters
    config = Configuration(
        range_shift_config=range_shift_config,
        cluster_link_config=cluster_link_config,
        iod_config=iod_config,
        od_config=od_config,
        odp_config=odp_config
    )

    orbits_completed = []
    continue_ = False
    if_exists_ = if_exists
    if out_dir is not None:
        if not os.path.exists(out_dir):
            os.mkdir(out_dir)
            logger.debug("Created {} directory.".format(out_dir))

        else:
            if if_exists == "continue":
                logger.warning("{} directory already exists, attempting to continue previous run.".format(out_dir))
                continue_ = True
            elif if_exists == "erase":
                logger.warning("{} directory already exists, removing previous results.".format(out_dir))
                shutil.rmtree(out_dir)
                os.mkdir(out_dir)
                logger.debug("Created {} directory.".format(out_dir))
            else:
                err = (
                    "if_exists should be one of {'continue', 'erase'}."
                )
                raise ValueError(err)

        # The primary files which will be used to determine if the run
        # can be continued from a previous state and, if so, from where
        # to continue the run
        config_file = os.path.join(out_dir, "config.yml")
        test_orbits_in_file = os.path.join(out_dir, "test_orbits_in.csv")
        status_file = os.path.join(out_dir, "status.txt")
        config_eq = False
        test_orbits_eq = False
        save_orbits = True
        save_config = True

        # Add summary file for test_orbits that tracks number of recovered orbits and number of observations
        # linked in addition to the test_orbit_id used by THOR
        test_orbits_out_file = os.path.join(out_dir, "test_orbits_out.csv")
        if continue_:
            if not os.path.exists(config_file):
                logger.warning("No previous configuration file found.")
                save_config = True
                if_exists_ = "erase"
            else:
                logger.info("Previous configuration file found. Comparing settings...")

                config_prev = Configuration.fromYaml(config_file)
                if config_prev != config:
                    logger.warning("Previous configuration does not match current configuration. Processing will not continue from previous state.")
                    if_exists_ = "erase"
                else:
                    config_eq = True
                    save_config = False
                    logger.info("Previous configuration matches current configuration.")

            if not os.path.exists(test_orbits_in_file):
                logger.warning("No previous test orbits file found.")
                save_orbits = True
            else:
                logger.info("Previous test orbits file found.")

                test_orbits_prev = Orbits.from_csv(
                    test_orbits_in_file,
                )
                if test_orbits_prev != test_orbits:
                    logger.warning("Previous test orbits do not match current test orbits.")
                else:
                    test_orbits_eq = True
                    save_orbits = False
                    test_orbits_df = test_orbits_prev.to_df(include_units=False)
                    logger.info("Previous test orbits match current test orbits.")

            if not os.path.exists(status_file):
                logger.warning("No previous status file found.")
            else:
                if test_orbits_eq and config_eq:
                    orbits_completed = np.loadtxt(
                        os.path.join(out_dir, "status.txt"),
                        delimiter="\n",
                        dtype=str,
                        ndmin=1
                    )
                    logger.info("Previous status file found.")

        if (not test_orbits_eq or not config_eq) and continue_:
            if if_exists == "continue":
                logger.critical("Previous run cannot continue from previous state.")
                raise ValueError("Previous run cannot continue from previous state. Set if_exists to 'erase' or change/delete the output directory.")
            elif if_exists == "erase":
                shutil.rmtree(out_dir)
                os.mkdir(out_dir)
                logger.debug("Created {} directory.".format(out_dir))
            else:
                pass

        if save_config:
            config.toYaml(config_file)
            logger.debug("Saved config.yml.")

        if save_orbits:
            test_orbits.to_csv(
                test_orbits_in_file
            )
            logger.debug("Saved test_orbits_in.csv.")

            preprocessed_observations.to_csv(
                os.path.join(out_dir, "preprocessed_observations.csv"),
                index=False,
                float_format="%.15e"
            )
            logger.debug("Saved preprocessed_observations.csv.")

    test_orbit_dfs = []
    recovered_orbits_dfs = []
    recovered_orbit_members_dfs = []
    obs_ids_linked = []
    num_orbits = len(test_orbits)
    if num_orbits != len(orbits_completed):
        test_orbits_split = test_orbits[len(orbits_completed):].split(1)
    else:
        test_orbits_split = []

    # If orbits have previously completed, read the results and continue iterating
    # through orbits not previously completed.
    id_offset = 0
    if len(orbits_completed) > 0:
        logger.info("{}/{} orbits have previously finished processing.".format(len(orbits_completed), num_orbits))

        test_orbits_df = Orbits.from_csv(
            test_orbits_out_file,
        ).to_df(include_units=False)
        logger.debug("Read previous test_orbits_out.csv.")

        recovered_orbits = Orbits.from_csv(
            os.path.join(out_dir, "recovered_orbits.csv"),
        ).to_df(include_units=False)
        logger.debug("Read previous recovered_orbits.csv.")

        recovered_orbit_members = pd.read_csv(
            os.path.join(out_dir, "recovered_orbit_members.csv"),
            index_col=False,
            dtype={"obs_id" : str},
            float_precision="round_trip"
        )
        logger.debug("Read previous recovered_orbit_members.csv.")

        test_orbit_dfs = [test_orbits_df]
        recovered_orbits_dfs = [recovered_orbits]
        recovered_orbit_members_dfs = [recovered_orbit_members]
        obs_ids_linked = recovered_orbit_members["obs_id"].values
        id_offset = len(orbits_completed)

    if len(test_orbits_split) != 0:
        for i, orbit_i in enumerate(test_orbits_split):

            time_start = time.time()
            orbit_id = "{:08d}".format(i + id_offset)

            logger.info("Processing orbit {} ({}/{})...".format(orbit_id, i + 1 + id_offset, num_orbits))

            if out_dir is not None:
                orbit_dir = os.path.join(out_dir, "orbit_{}".format(orbit_id))
            else:
                orbit_dir = None

            linked_mask = (~preprocessed_observations["obs_id"].isin(obs_ids_linked))

            recovered_orbits_i, recovered_orbit_members_i = runTHOROrbit(
                preprocessed_observations[linked_mask],
                orbit_i,
                range_shift_config=range_shift_config,
                cluster_link_config=cluster_link_config,
                iod_config=iod_config,
                od_config=od_config,
                odp_config=odp_config,
                out_dir=orbit_dir,
                if_exists=if_exists_,
                logging_level=logging_level
            )

            time_end = time.time()

            if len(recovered_orbits_i) > 0:
                recovered_orbits_i.insert(0, "test_orbit_id", orbit_id)
                recovered_orbit_members_i.insert(0, "test_orbit_id", orbit_id)
                obs_ids_linked_i = recovered_orbit_members_i["obs_id"].unique()
                obs_ids_linked = np.concatenate([obs_ids_linked, obs_ids_linked_i])

                orbits_recovered = len(recovered_orbits_i)
                observations_linked = len(obs_ids_linked_i)
            else:
                orbits_recovered = 0
                observations_linked = 0

            test_orbit_i = orbit_i.to_df(include_units=False)
            test_orbit_i["test_orbit_id"] = orbit_id
            test_orbit_i["orbits_recovered"] = orbits_recovered
            test_orbit_i["observations_linked"] = observations_linked
            test_orbit_i["processing_time"] = time_end - time_start
            test_orbit_dfs.append(test_orbit_i)

            logger.info("Completed processing orbit {} in {:.3f} seconds.".format(orbit_id, time_end - time_start))

            recovered_orbits_dfs.append(recovered_orbits_i)
            recovered_orbit_members_dfs.append(recovered_orbit_members_i)

            test_orbits_df = pd.concat(
                test_orbit_dfs,
                ignore_index=True
            )
            recovered_orbits = pd.concat(
                recovered_orbits_dfs,
                ignore_index=True
            )
            recovered_orbit_members = pd.concat(
                recovered_orbit_members_dfs,
                ignore_index=True
            )

            if out_dir is not None:
                Orbits.from_df(test_orbits_df).to_csv(
                    test_orbits_out_file
                )
                logger.debug("Saved test_orbits_out.csv.")

                Orbits.from_df(recovered_orbits).to_csv(
                    os.path.join(out_dir, "recovered_orbits.csv")
                )
                logger.debug("Saved recovered_orbits.csv.")

                recovered_orbit_members.to_csv(
                    os.path.join(out_dir, "recovered_orbit_members.csv"),
                    index=False,
                    float_format="%.15e"
                )
                logger.debug("Saved recovered_orbit_members.csv.")

            orbits_completed = np.concatenate([orbits_completed, np.array([orbit_id])])
            if out_dir is not None:
                with open(os.path.join(out_dir, "status.txt"), "w") as status_out:
                    np.savetxt(
                        status_out,
                        orbits_completed,
                        delimiter="\n",
                        fmt="%s"
                    )
                logger.info("Saved status.txt.")

    else:

        logger.info("Run completed previously.")
        test_orbits_df = Orbits.from_csv(
                test_orbits_out_file,
        ).to_df(include_units=False)
        logger.debug("Read test_orbits_out.csv.")

        recovered_orbits = Orbits.from_csv(
            os.path.join(out_dir, "recovered_orbits.csv"),
        ).to_df(include_units=False)
        logger.debug("Read recovered_orbits.csv.")

        recovered_orbit_members = pd.read_csv(
            os.path.join(out_dir, "recovered_orbit_members.csv"),
            index_col=False,
            dtype={"obs_id" : str},
            float_precision="round_trip"
        )
        logger.debug("Read recovered_orbit_members.csv.")

    return test_orbits_df, recovered_orbits, recovered_orbit_members


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#    Project: Fast Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2017-2018 European Synchrotron Radiation Facility, Grenoble, France
#
#    Principal author:       Jérôme Kieffer (Jerome.Kieffer@ESRF.eu)
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#  .
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#  .
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#  THE SOFTWARE.

"""
Detectors manufactured by PSI, those may be different from the one from Dectris
"""

__author__ = "Jerome Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "2021 European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "15/03/2021"
__status__ = "production"

import numpy
import logging
from ._common import Detector
from ..utils import mathutil
logger = logging.getLogger(__name__)

class Jungfrau(Detector):
    """
    Raw Jungfrau module without sub-module pixel expension applied.
    """
    MANUFACTURER = "PSI"

    MODULE_SIZE = (256, 256)  # number of pixels per module (y, x)
    MAX_SHAPE = (512, 1024)  # max size of the detector
    PIXEL_SIZE = (75e-6, 75e-6)
    BORDER_SIZE_RELATIVE = 2.0
    force_pixel = True
    aliases = ["Jungfrau 500k"]
    uniform_pixel = False

    @classmethod
    def _calc_pixels_size(cls, length, module_size, pixel_size):
        """
        given the length (in pixel) of the detector, the size of a
        module (in pixels) and the pixel_size (in meter). this method
        return the length of each pixels 0..length.

        :param length: the number of pixel to compute
        :type length: int
        :param module_size: the number of pixel of one module
        :type module_size: int
        :param pixel_size: the size of one pixels (meter per pixel)
        :type length: float

        :return: the coordinates of each pixels 0..length
        :rtype: ndarray
        """
        size = numpy.ones(length)
        n = length // module_size
        for i in range(1, n):
            size[i * module_size - 1] = cls.BORDER_SIZE_RELATIVE
            size[i * module_size] = cls.BORDER_SIZE_RELATIVE
        return pixel_size * size

    def __init__(self, pixel1=75e-6, pixel2=75e-6, max_shape=None, module_size=None):
        Detector.__init__(self, pixel1=pixel1, pixel2=pixel2, max_shape=max_shape)
        self._pixel_edges = None  # array of size max_shape+1: pixels are contiguous
        if (module_size is None) and ("MODULE_SIZE" in dir(self.__class__)):
            self.module_size = tuple(self.MODULE_SIZE)
        else:
            self.module_size = module_size

    def __repr__(self):
        return "Detector %s\t PixelSize= %.3e, %.3e m" % \
            (self.name, self.pixel1, self.pixel2)

    def calc_pixels_edges(self):
        """
        Calculate the position of the pixel edges
        """
        if self._pixel_edges is None:
            pixel_size1 = self._calc_pixels_size(self.max_shape[0], self.module_size[0], self.PIXEL_SIZE[0])
            pixel_size2 = self._calc_pixels_size(self.max_shape[1], self.module_size[1], self.PIXEL_SIZE[1])
            pixel_edges1 = numpy.zeros(self.max_shape[0] + 1)
            pixel_edges2 = numpy.zeros(self.max_shape[1] + 1)
            pixel_edges1[1:] = numpy.cumsum(pixel_size1)
            pixel_edges2[1:] = numpy.cumsum(pixel_size2)
            self._pixel_edges = pixel_edges1, pixel_edges2
        return self._pixel_edges

    def get_pixel_corners(self, correct_binning=False):
        """
        Calculate the position of the corner of the pixels

        This should be overwritten by class representing non-contiguous detector (Xpad, ...)

        Precision float32 is ok: precision of 1µm for a detector size of 1m

        :return:  4D array containing:
                    pixel index (slow dimension)
                    pixel index (fast dimension)
                    corner index (A, B, C or D), triangles or hexagons can be handled the same way
                    vertex position (z,y,x)
        """

        if self._pixel_corners is None:
            with self._sem:
                if self._pixel_corners is None:
                    edges1, edges2 = self.calc_pixels_edges()
                    p1 = mathutil.expand2d(edges1, self.shape[1] + 1, False)
                    p2 = mathutil.expand2d(edges2, self.shape[0] + 1, True)
                    # p3 = None
                    self._pixel_corners = numpy.zeros((self.shape[0], self.shape[1], 4, 3), dtype=numpy.float32)
                    self._pixel_corners[:,:, 0, 1] = p1[:-1,:-1]
                    self._pixel_corners[:,:, 0, 2] = p2[:-1,:-1]
                    self._pixel_corners[:,:, 1, 1] = p1[1:,:-1]
                    self._pixel_corners[:,:, 1, 2] = p2[1:,:-1]
                    self._pixel_corners[:,:, 2, 1] = p1[1:, 1:]
                    self._pixel_corners[:,:, 2, 2] = p2[1:, 1:]
                    self._pixel_corners[:,:, 3, 1] = p1[:-1, 1:]
                    self._pixel_corners[:,:, 3, 2] = p2[:-1, 1:]
                    # if p3 is not None:
                    #     # non flat detector
                    #    self._pixel_corners[:, :, 0, 0] = p3[:-1, :-1]
                    #     self._pixel_corners[:, :, 1, 0] = p3[1:, :-1]
                    #     self._pixel_corners[:, :, 2, 0] = p3[1:, 1:]
                    #     self._pixel_corners[:, :, 3, 0] = p3[:-1, 1:]
        if correct_binning and self._pixel_corners.shape[:2] != self.shape:
            return self._rebin_pixel_corners()
        else:
            return self._pixel_corners

    def calc_cartesian_positions(self, d1=None, d2=None, center=True, use_cython=True):
        """
        Calculate the position of each pixel center in cartesian coordinate
        and in meter of a couple of coordinates.
        The half pixel offset is taken into account here !!!

        :param d1: the Y pixel positions (slow dimension)
        :type d1: ndarray (1D or 2D)
        :param d2: the X pixel positions (fast dimension)
        :type d2: ndarray (1D or 2D)

        :return: position in meter of the center of each pixels.
        :rtype: ndarray

        d1 and d2 must have the same shape, returned array will have
        the same shape.

        """
        # if the detctor has been tweaked with an ASCII geometry ... fall-back on the classical method:
        if self._pixel_corners is not None:
            return Detector.calc_cartesian_positions(self, d1=d1, d2=d2, center=center, use_cython=use_cython)

        edges1, edges2 = self.calc_pixels_edges()

        if (d1 is None) or (d2 is None):
            if center:
                # Take the center of each pixel
                d1 = 0.5 * (edges1[:-1] + edges1[1:])
                d2 = 0.5 * (edges2[:-1] + edges2[1:])
            else:
                # take the lower corner
                d1 = edges1[:-1]
                d2 = edges2[:-1]
            p1 = numpy.outer(d1, numpy.ones(self.shape[1]))
            p2 = numpy.outer(numpy.ones(self.shape[0]), d2)
        else:
            if center:
                # Not +=: do not mangle in place arrays
                d1 = d1 + 0.5
                d2 = d2 + 0.5
            p1 = numpy.interp(d1, numpy.arange(self.max_shape[0] + 1), edges1, edges1[0], edges1[-1])
            p2 = numpy.interp(d2, numpy.arange(self.max_shape[1] + 1), edges2, edges2[0], edges2[-1])
        return p1, p2, None

class Jungfrau_16M_cor(Jungfrau):
    """Jungfrau 16 corrected for double-sized pixels
    """
    MODULE_SIZE = ((512 + 2), 1024 + 6)  # number of pixels per module (y, x)
    MAX_SHAPE = ((512 + 2) * 32, 1024 + 6)  # max size of the detector
    force_pixel = True
    aliases = ["Jungfrau 16M cor"]

    @staticmethod
    def load_geom(geom_fname):
        """"Load module geometry from ASCII file

        Stollen from Alejandro Homs' code
        """
        import re
        geom_re = re.compile('m(?P<mod>[0-9]+)/(?P<par>[^ \t=]+)[ \t]*='
                         '[ \t]*(?P<val>.+)')
        module_geom = {}
        with open(geom_fname) as ifile:
            for l in ifile:
                m = geom_re.match(l)
                if not m:
                    continue
                mod = int(m.group('mod'))
                mod_data = module_geom.setdefault(mod, {})
                val = m.group('val')
                if ' ' in val:
                    val = val.split()
                    if val[0].endswith('x') and val[1].endswith('y'):
                        val = [v[:-1] for v in val]
                else:
                    val = [val]
                key = m.group('par')
                if key.startswith("min_") or key.startswith("max_"):
                    mod_data[key] = int(val[0])
                elif key.startswith("corner"):
                    mod_data[key] = float(val[0])
                else:
                    mod_data[key] = [float(v) for v in val]
        return module_geom

    def init_from_geometry(self, filename):
        """initialize the detector from "geom" file produced at  PSI"""
        config = self.load_geom(filename)
        shape0 = 0
        shape1 = 0
        for m in config.values():
            shape0 = max(shape0, m.get("max_ss", 0))
            shape1 = max(shape1, m.get("max_fs", 0))
        self.MAX_SHAPE = (shape0 + 1, shape1 + 1)

        position_array = numpy.zeros(self.MAX_SHAPE + (4, 3), dtype=numpy.float32)

        for module in config.values():
            slab = position_array[module["min_ss"]: 1 + module["max_ss"], module["min_fs"]: 1 + module["max_fs"]]
            ss_edges = numpy.arange(2 + module["max_ss"] - module["min_ss"], dtype=numpy.int32)
            fs_edges = numpy.arange(2 + module["max_fs"] - module["min_fs"], dtype=numpy.int32)
            p1 = mathutil.expand2d(ss_edges, fs_edges.size, False)
            p2 = mathutil.expand2d(fs_edges, ss_edges.size, True)
            indexes = numpy.vstack([p2.ravel(), p1.ravel()])  # XY
            mat = numpy.array([module["fs"], module["ss"]], dtype=numpy.float64)
            position_xy = mat.dot(indexes) + numpy.array([[module["corner_x"]], [module["corner_y"]]])
            p2, p1 = position_xy.reshape((2,) + p1.shape)
            slab[:,:, 0, 1] = p1[:-1,:-1]
            slab[:,:, 0, 2] = p2[:-1,:-1]
            slab[:,:, 1, 1] = p1[1:,:-1]
            slab[:,:, 1, 2] = p2[1:,:-1]
            slab[:,:, 2, 1] = p1[1:, 1:]
            slab[:,:, 2, 2] = p2[1:, 1:]
            slab[:,:, 3, 1] = p1[:-1, 1:]
            slab[:,:, 3, 2] = p2[:-1, 1:]
        self._pixel_corners = (position_array * self.pixel1).astype(numpy.float32)
        self.IS_CONTIGUOUS = False

    def calc_mask(self):
        "Mask out sub-module junctions"
        mask = numpy.zeros(self.MAX_SHAPE, dtype=numpy.int8)
        mask[255:self.MAX_SHAPE[0] - 2:self.MODULE_SIZE[0]] = 1
        mask[256:self.MAX_SHAPE[0] - 2:self.MODULE_SIZE[0]] = 1
        mask[257:self.MAX_SHAPE[0] - 2:self.MODULE_SIZE[0]] = 1
        mask[258:self.MAX_SHAPE[0] - 2:self.MODULE_SIZE[0]] = 1

        for i in range(0, self.MODULE_SIZE[1], 258):
            mask[:, i + 255:self.MAX_SHAPE[1] - 2:self.MODULE_SIZE[1]] = 1
            mask[:, i + 256:self.MAX_SHAPE[1] - 2:self.MODULE_SIZE[1]] = 1
            mask[:, i + 257:self.MAX_SHAPE[1] - 2:self.MODULE_SIZE[1]] = 1
            mask[:, i + 258:self.MAX_SHAPE[1] - 2:self.MODULE_SIZE[1]] = 1
        return mask


"""
This module contains the classes and functions to generate the zmatrix of
a molecule.
"""

from copy import deepcopy

import numpy as np

from peleffy.topology.elements import DummyAtom

class ZMatrix(np.ndarray):
    """
    It generates the zmatrix of a molecule as a numpy.array.

    Inspired by the PlopRotTemp algorithm.
    """

    def __init__(self, topology):
        """
        It initializes a ZMatrix object.

        Parameters
        ----------
        topology : a peleffy.topology.Topology
            The molecular topology representation to generate the
            zmatrix with

        Examples
        --------

        Given a molecular topology, build its Z-matrix

        >>> from peleffy.topology import Molecule

        >>> molecule = Molecule(smiles='Cc1ccccc1')

        >>> from peleffy.forcefield import OpenForceField

        >>> openff = OpenForceField('openff_unconstrained-1.2.1.offxml')
        >>> parameters = openff.parameterize(molecule)

        >>> from peleffy.topology import Topology
        >>> topology = Topology(molecule, parameters)

        >>> from peleffy.topology import ZMatrix

        >>> zmatrix = ZMatrix(topology)

        """
        # We will work on a copy of the topology's object to modify it freely
        self._topology = deepcopy(topology)

    def __new__(cls, topology):
        """
        It customizes the creation of the numpy.array.

        Parameters
        ----------
        topology : a peleffy.topology.Topology
            The molecular topology representation to generate the
            zmatrix with
        """
        topology = deepcopy(topology)
        obj = np.zeros((len(topology.atoms), 3)).view(cls)
        coords = cls._extract_coords(topology)
        obj = cls._build_zmatrix(cls, obj, coords, topology)

        return obj

    @staticmethod
    def _extract_coords(topology):
        """
        It extracts the coordinates of the molecule's atoms.

        Parameters
        ----------
        topology : a peleffy.topology.Topology
            The molecular topology representation to generate the
            zmatrix with

        Returns
        -------
        coords : list[tuple[float]]
            The coordinates of the molecule
        """
        coords = list()
        for atom in topology.atoms:
            coords.append((atom.x, atom.y, atom.z))

        return coords

    @staticmethod
    def _get_absolute_parent(topology):
        """
        It returns the absolute parent in the topology of the molecule.

        Parameters
        ----------
        topology : a peleffy.topology.Topology
            The molecular topology representation to generate the
            zmatrix with

        Returns
        -------
        absolute_parent : peleffy.topology.molecule.Atom
            The absolute parent of the molecule
        """
        absolute_parent = list()
        for atom in topology.atoms:
            if atom.parent is None:
                absolute_parent.append(atom)

        assert len(absolute_parent) == 1, 'Only 1 absolute parent is expected'

        return absolute_parent[0]

    @staticmethod
    def _calculate_bond(x1, y1, z1, x2, y2, z2):
        """
        It calculates the bond distance between two sets of coordinates.

        Parameters
        ----------
        x1 : float
            The x1 coordinate
        y1 : float
            The y1 coordinate
        z1 : float
            The z1 coordinate
        x2 : float
            The x2 coordinate
        y2 : float
            The y2 coordinate
        z2 : float
            The z2 coordinate

        Returns
        -------
        distance : float
            The bond distance between the two sets of coordinates
        """
        dx = x1 - x2
        dy = y1 - y2
        dz = z1 - z2

        return np.sqrt(dx * dx + dy * dy + dz * dz)

    @staticmethod
    def _calculate_angle(x1, y1, z1, x2, y2, z2, x3, y3, z3):
        """
        It calculates the angle between three sets of coordinates.

        Parameters
        ----------
        x1 : float
            The x1 coordinate
        y1 : float
            The y1 coordinate
        z1 : float
            The z1 coordinate
        x2 : float
            The x2 coordinate
        y2 : float
            The y2 coordinate
        z2 : float
            The z2 coordinate
        x3 : float
            The x3 coordinate
        y3 : float
            The y3 coordinate
        z3 : float
            The z3 coordinate

        Returns
        -------
        angle : float
            The angle between the three sets of coordinates
        """
        dx_12 = x1 - x2
        dy_12 = y1 - y2
        dz_12 = z1 - z2

        dx_31 = x3 - x1
        dy_31 = y3 - y1
        dz_31 = z3 - z1

        vdot = dx_12 * dx_31 + dy_12 * dy_31 + dz_12 * dz_31

        assert np.fabs(vdot) > 0, 'A non-zero angle is expected'

        d12 = np.sqrt(dx_12 * dx_12 + dy_12 * dy_12 + dz_12 * dz_12)
        d31 = np.sqrt(dx_31 * dx_31 + dy_31 * dy_31 + dz_31 * dz_31)

        xang = vdot / (d12 * d31)

        if xang - 1.0 > -0.0000000001:
            return 0.0
        elif xang + 1.0 < 0.0000000001:
            return np.pi

        return np.arccos(xang) * 180.0 / np.pi

    @staticmethod
    def _calculate_dihedral(x1, y1, z1, x2, y2, z2, x3, y3, z3, x4, y4, z4):
        """
        It calculates the dihedral between four sets of coordinates.

        Parameters
        ----------
        x1 : float
            The x1 coordinate
        y1 : float
            The y1 coordinate
        z1 : float
            The z1 coordinate
        x2 : float
            The x2 coordinate
        y2 : float
            The y2 coordinate
        z2 : float
            The z2 coordinate
        x3 : float
            The x3 coordinate
        y3 : float
            The y3 coordinate
        z3 : float
            The z3 coordinate
        x4 : float
            The x4 coordinate
        y4 : float
            The y4 coordinate
        z4 : float
            The z4 coordinate

        Returns
        -------
        dihedral : float
            The dihedral between the four sets of coordinates
        """
        dx_12 = x1 - x2
        dy_12 = y1 - y2
        dz_12 = z1 - z2

        dx_32 = x3 - x2
        dy_32 = y3 - y2
        dz_32 = z3 - z2

        dx_34 = x3 - x4
        dy_34 = y3 - y4
        dz_34 = z3 - z4

        ax = dy_12 * dz_32 - dz_12 * dy_32
        ay = dz_12 * dx_32 - dx_12 * dz_32
        az = dx_12 * dy_32 - dy_12 * dx_32
        cx = dy_32 * dz_34 - dz_32 * dy_34
        cy = dz_32 * dx_34 - dx_32 * dz_34
        cz = dx_32 * dy_34 - dy_32 * dx_34

        rac = ax * cx + ay * cy + az * cz
        ra = ax * ax + ay * ay + az * az
        rc = cx * cx + cy * cy + cz * cz

        cosang = rac / np.sqrt(ra * rc)

        if cosang - 1.0 > -0.00000000001:
            phi = 0.0
        elif cosang + 1.0 < 0.00000000001:
            phi = np.pi
        else:
            phi = np.arccos(cosang)

        s = dx_12 * cx + dy_12 * cy + dz_12 * cz
        if (s < 0):
            phi = -phi  # to account for phi between pi and 2pi

        return phi * 180.0 / np.pi

    @staticmethod
    def _build_zmatrix(cls, obj, coords, topology):
        """
        It buils the zmatrix.

        Parameters
        ----------
        cls : ZMatrix class
            The ZMatrix class
        obj : ZMatrix object
            The ZMatrix object
        coords : list[tuple[float]]
            The coordinates of the molecule
        topology : a peleffy.topology.Topology
            The molecular topology representation to generate the
            zmatrix with

        Returns
        -------
        obj :  ZMatrix object
            The ZMatrix object which is a numpy.array with the corresponding
            zmatrix initialized.
        """
        dummy1 = DummyAtom(index=-3, PDB_name='DUM1', parent=None)
        dummy2 = DummyAtom(index=-2, PDB_name='DUM2', parent=dummy1)
        dummy3 = DummyAtom(index=-1, PDB_name='DUM3', parent=dummy2)
        absolute_parent = cls._get_absolute_parent(topology)
        absolute_parent.set_parent(dummy3)

        dummy1.set_coords([0.0, 0.0, 0.0])  # at origin
        dummy2.set_coords([0.0, 0.0, 1.0])  # at z axis
        dummy3.set_coords([1.0, 0.0, 0.0])  # at x axis

        for i, atom in enumerate(topology.atoms):
            atom1 = atom
            atom2 = atom1.parent
            atom3 = atom2.parent
            atom4 = atom3.parent

            assert atom1 is not None and atom2 is not None \
                and atom3 is not None and atom4 is not None, \
                'A None as parent is not expected'

            x1, y1, z1 = (atom1.x, atom1.y, atom1.z)
            x2, y2, z2 = (atom2.x, atom2.y, atom2.z)
            x3, y3, z3 = (atom3.x, atom3.y, atom3.z)
            x4, y4, z4 = (atom4.x, atom4.y, atom4.z)

            obj[i][0] = cls._calculate_bond(x1, y1, z1,
                                            x2, y2, z2)
            obj[i][1] = cls._calculate_angle(x1, y1, z1,
                                             x2, y2, z2,
                                             x3, y3, z3)
            obj[i][2] = cls._calculate_dihedral(x1, y1, z1,
                                                x2, y2, z2,
                                                x3, y3, z3,
                                                x4, y4, z4)

        return obj


"""
A GP prior over qso spectrum for zQSO estimation.
"""
from .null_gp import NullGP

import numpy as np
from scipy import interpolate
import h5py

from .zqso_set_parameters import ZParameters
from .zqso_samples import ZSamples

class ZGP(NullGP):
    """
    GP inference on zQSO:
        p(y | λ, σ², M, ω, blue_σ², red_σ²),
    where we also model the data outside of the modelling window
    as i.i.d Gaussians. See details in
    <Automated Measurement of Quasar Redshift with a Gaussian Process
    https://arxiv.org/abs/2006.07343>

    :param rest_wavelengths: λ, the range of λ you model your GP on QSO emission
    :param mu: mu, the mean model of the GP.
    :param M: M, the low rank decomposition of the covariance kernel: K = MM^T.
    :param bluewards_mu: the mean model for blueward part of the data outside of
        the modelling window.
    :param bluewards_sigma: the noise model for blueward part of the data outside of
        the modelling window.
    :param redwards_mu: the mean model for redward part of the data outside of
        the modelling window.
    :param redwards_sigma: the noise model for redward part of the data outside of
        the modelling window.
    """

    def __init__(
        self,
        params: ZParameters,
        z_qso_samples: ZSamples,
        rest_wavelengths: np.ndarray,
        mu: np.ndarray,
        M: np.ndarray,
        bluewards_mu: float,
        redwards_mu: float,
        bluewards_sigma: float,
        redwards_sigma: float,
    ):
        self.params = params
        self.z_qso_samples = z_qso_samples

        # learned GP model
        self.rest_wavelengths = rest_wavelengths
        self.mu = mu
        self.M = M
        self.bluewards_mu = bluewards_mu
        self.redwards_mu = redwards_mu
        self.bluewards_sigma = bluewards_sigma
        self.redwards_sigma = redwards_sigma

        # preprocess model interpolants
        self.mu_interpolator = interpolate.interp1d(rest_wavelengths, mu)
        self.list_M_interpolators = [
            interpolate.interp1d(rest_wavelengths, eigenvector) for eigenvector in M.T
        ]

    def get_interp(
        self, x: np.ndarray, y: np.ndarray, wavelengths: np.ndarray, z_qso: float
    ) -> None:
        """
        Build and interpolate the GP model onto the observed data.

        p(y | λ, zqso, v, ω, M_nodla) = N(y; μ, (K + Ω) + V)

        :param x: this_rest_wavelengths
        :param y: this_flux
        :param wavelengths: observed wavelengths, put this arg is just for making the
            code looks similar to the MATLAB code. This could be taken off in the future.
        :param z_qso: quasar redshift

        Note: assume already pixel masked.
        """
        # interpolate model onto given wavelengths
        this_mu = self.mu_interpolator(x)
        this_M = self.M_interpolator(x)

        assert this_M.shape[1] == self.params.k

        # assign to instance attrs
        self.this_mu = this_mu
        self.this_M = this_M

    def set_data(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        noise_variance: np.ndarray,
        pixel_mask: np.ndarray,
        z_qso: float,
        normalize: bool = True,
        build_model: bool = True,
    ) -> None:
        """
        Set "testing" data to be evaluated. Now assumed to be a single
        spectrum, but we can implement a batch of spectra in the future.

        The difference between this and the null model set_data is that
        we need to keep all of the data no matter they are inside or
        outside the modelling range, and extend the model to fill up
        all of the data points.

        :param X: (n_points, ) this_wavelengths, the `OBSERVED` wavelengths of the spectrum.
        :param Y: (n_points, ) this_flux, the flux of the observed spectrum
        :param noise_variance: (n_points, ) the instrumental noise variance per pixel.
        :param pixel_mask: (n_points, ) the pixel mask corresponding to the read_spec you used.
        :param z_qso: the redshift of the quasar.

        Note: n_points represents number of pixels within a spectrum
        """
        # variant redshift in quasars
        self.z_qso = z_qso

        # cut-off observations
        max_pos_lambda = self.params.observed_wavelengths(self.params.max_lambda, z_qso)
        min_pos_lambda = self.params.observed_wavelengths(self.params.min_lambda, z_qso)

        max_observed_lambda = np.min((max_pos_lambda, np.max(X)))
        min_observed_lambda = np.max((min_pos_lambda, np.min(X)))

        labmda_observed = max_observed_lambda - min_observed_lambda

        # filter the spectrum; these quantities are within the modelling window
        ind = (X > min_observed_lambda) * (X < max_observed_lambda)
        self.y = Y[ind]
        self.this_wavelengths = X[ind]
        self.v = noise_variance[ind]
        self.pixel_mask = pixel_mask[ind]

        # convert to QSO rest frame
        self.x = self.params.emitted_wavelengths(X[ind], z_qso)

        # normalize flux
        if normalize:
            ind = (self.x >= self.params.normalization_min_lambda) & (
                self.x <= self.params.normalization_max_lambda
            )
            this_median = np.nanmedian(self.y[ind])
            self.y = self.y / this_median
            self.v = self.v / this_median ** 2

        # Normalise the observed flux for out-of-range model since the
        # redward- and blueward- models were trained with normalisation.
        # Find probability for out-of-range model
        this_normalized_flux = (
            Y / this_median
        )  # since we've modified self.y, we thus need to
        # use this_out_flux which is outside the parfor loop
        this_normalized_v = noise_variance / this_median ** 2

        # select blueward region
        ind_bw = (X < min_observed_lambda) & (~pixel_mask)
        self.y_bw = this_normalized_flux[ind_bw]
        self.v_bw = this_normalized_v[ind_bw]
        # select redward region
        ind_rw = (X > max_observed_lambda) & (~pixel_mask)
        self.y_rw = this_normalized_flux[ind_rw]
        self.v_rw = this_normalized_v[ind_rw]

        # apply pixel mask and filter spectrum within modelling range
        ind = (self.x >= self.params.min_lambda) & (self.x <= self.params.max_lambda)
        ind = ind & (~self.pixel_mask)

        self.this_wavelengths = self.this_wavelengths[ind]
        self.x = self.x[ind]
        self.y = self.y[ind]
        self.v = self.v[ind]

        self.v[np.isinf(self.v)] = np.nanmean(self.v)  # rare kludge to fix bad data

        self.ind = ind

        if build_model:
            self.get_interp(self.x, self.y, self.this_wavelengths, self.z_qso)

    def log_model_evidence(self) -> float:
        """
        Compute the log model evidence with the learned model.

        Note: assume the model has already interpolated onto the observed data,
        and data got loaded into the GP instance.
        """
        # log likelihood within the modelling window
        log_likelihood = self.log_mvnpdf_low_rank(
            self.y, self.this_mu, self.this_M, self.v
        )

        n_bw = self.y_bw.shape[0]
        n_rw = self.y_rw.shape[0]

        # calculate log likelihood of iid multivariate normal with
        #   log N(y; mu, diag(V) + sigma^2 )
        bw_log_likelihood = self.log_mvnpdf_iid(
            self.y_bw,
            self.bluewards_mu * np.ones((n_bw,)),
            self.bluewards_sigma ** 2 * np.ones((n_bw,)) + self.v_bw,
        )
        rw_log_likelihood = self.log_mvnpdf_iid(
            self.y_rw,
            self.redwards_mu * np.ones((n_rw,)),
            self.redwards_sigma ** 2 * np.ones((n_rw,)) + self.v_rw,
        )

        return log_likelihood + bw_log_likelihood + rw_log_likelihood

    def inference_z_qso(
        self,
        wavelengths: np.ndarray,
        flux: np.ndarray,
        noise_variance: np.ndarray,
        pixel_mask: np.ndarray,
        z_qso_min: float = 2.14,
        z_qso_max: float = 6.16,
    ):
        """
        Sample the zQSO within a prior volume define in self.z_qso_samples
        """
        sample_log_likelihoods = np.full((self.z_qso_samples.num_zqso_samples,), np.nan)
        sample_z_qsos = self.z_qso_samples.sample_z_qsos(
            z_qso_min=z_qso_min, z_qso_max=z_qso_max
        )

        for i, z_qso in enumerate(sample_z_qsos):
            # set the data and interpolate the model
            self.set_data(
                wavelengths,
                flux,
                noise_variance,
                pixel_mask,
                z_qso=z_qso,
                normalize=True,
                build_model=True,
            )

            sample_log_likelihoods[i] = self.log_model_evidence()

        self.sample_log_likelihoods = sample_log_likelihoods

        # maximum a posteriori
        I = np.nanargmax(sample_log_likelihoods)
        self.z_map = sample_z_qsos[I]
        print("[Info] Z MAP = {:.3g}".format(self.z_map))

    @staticmethod
    def log_mvnpdf_iid(y: np.ndarray, mu: np.ndarray, d: np.ndarray,) -> float:
        """
        computes mutlivariate normal dist with
        each dim is iid, so no covariance.
            log N(y; mu, diag(d))

        :param y: this_flux, (n_points, )
        :param mu: this_mu, the mean vector of GP, (n_points, )
        :param d: diagonal noise term, (n_points, )
        """
        log_2pi = 1.83787706640934534

        n = d.shape[0]

        y = y[:, None] - mu[:, None]

        d_inv = 1 / d[:, None]  # (n_points, 1)
        D_inv_y = d_inv * y  # (n_points, 1)

        K_inv_y = D_inv_y  # (n_points, 1)

        log_det_K = np.sum(np.log(d))

        log_p = -0.5 * (np.matmul(y.T, K_inv_y).sum() + log_det_K + n * log_2pi)

        return log_p

    @property
    def this_noise(self):
        """
        noise kernel: instrumental noise
        """
        return self.v

class ZGPMAT(ZGP):
    """
    Load learned model from .mat file
    """

    def __init__(
        self,
        params: ZParameters,
        z_qso_samples: ZSamples,
        learned_file: str = "learned_zqso_only_model_outdata_full_dr9q_minus_concordance_norm_1176-1256.mat",
    ):
        with h5py.File(learned_file, "r") as learned:

            rest_wavelengths = learned["rest_wavelengths"][:, 0]
            mu = learned["mu"][:, 0]
            M = learned["M"][()].T
            bluewards_mu = learned["bluewards_mu"][0, 0]
            redwards_mu = learned["redwards_mu"][0, 0]
            bluewards_sigma = learned["bluewards_sigma"][0, 0]
            redwards_sigma = learned["redwards_sigma"][0, 0]

        super().__init__(
            params=params,
            z_qso_samples=z_qso_samples,
            rest_wavelengths=rest_wavelengths,
            mu=mu,
            M=M,
            bluewards_mu=bluewards_mu,
            redwards_mu=redwards_mu,
            bluewards_sigma=bluewards_sigma,
            redwards_sigma=redwards_sigma,
        )