#!/usr/bin/env python
#
# dials.ellipsoid.py
#
#  Copyright (C) 2018 Diamond Light Source
#
#  Author: James Parkhurst
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.

from __future__ import annotations

import copy
import logging
from math import pi

import numpy as np
from numpy.linalg import inv, norm

from libtbx.phil import parse

from dials.algorithms.profile_model.ellipsoid import calc_s1_s2, chisq_pdf
from dials.algorithms.profile_model.ellipsoid.model import ProfileModelFactory
from dials.algorithms.profile_model.ellipsoid.parameterisation import ModelState
from dials.algorithms.profile_model.ellipsoid.refiner import Refiner as ProfileRefiner
from dials.algorithms.profile_model.ellipsoid.refiner import RefinerData
from dials.algorithms.profile_model.gaussian_rs import BBoxCalculator, MaskCalculator
from dials.algorithms.profile_model.gaussian_rs.calculator import (
    ComputeEsdBeamDivergence,
)
from dials.algorithms.spot_prediction import IndexGenerator
from dials.array_family import flex
from dials.constants import FULL_PARTIALITY

logger = logging.getLogger("dials")

ellipsoid_algorithm_phil_scope = parse(
    """
  integration {

    use_crude_shoebox_mask = False
      .type = bool

    shoebox {

      probability = %f
        .type = float

    }

    corrections {
      lp = True
        .type = bool
      dqe = True
        .type = bool
      partiality = True
        .type = bool
    }

  }

"""
    % FULL_PARTIALITY
)

# Parameters
phil_scope = parse(
    """
include scope dials.algorithms.profile_model.ellipsoid.model.phil_scope
""",
    process_includes=True,
).adopt_scope(ellipsoid_algorithm_phil_scope)


def _compute_beam_vector(experiment, reflection_table):
    """Compute the obseved beam vector"""
    s1_obs = flex.vec3_double(len(reflection_table))
    for i, (panel_id, xyzobs) in enumerate(
        zip(reflection_table["panel"], reflection_table["xyzobs.px.value"])
    ):
        x, y, _ = xyzobs
        s1_obs[i] = experiment.detector[panel_id].get_pixel_lab_coord((x, y))
    return s1_obs


def _compute_bbox(experiment, reflection_table, sigma_d, s1="s1_obs"):
    """Compute the bounding box"""

    # Initialise the bounding box calculator
    compute_bbox = BBoxCalculator(
        experiment.crystal,
        experiment.beam,
        experiment.detector,
        experiment.goniometer,
        experiment.scan,
        sigma_d * 6,
        0,
    )

    # Compute the bounding box
    bbox = compute_bbox(
        reflection_table[s1],
        reflection_table["xyzcal.px"].parts()[2],
        reflection_table["panel"],
    )

    return bbox


def _compute_mask(
    experiment, reflection_table, sigma_d, s1="s1_obs", strong_shoeboxes=None
):
    """Compute the spot mask"""

    # Initialise the mask calculator
    mask_foreground = MaskCalculator(
        experiment.crystal,
        experiment.beam,
        experiment.detector,
        experiment.goniometer,
        experiment.scan,
        sigma_d * 3,
        0,
    )

    # Compute the reflection mask
    mask_foreground(
        reflection_table["shoebox"],
        reflection_table[s1],
        reflection_table["xyzcal.px"].parts()[2],
        reflection_table["panel"],
    )

    # if strong_shoeboxes given, apply the mask from this
    if strong_shoeboxes:
        # Apply strong spot mask
        assert len(reflection_table) == len(strong_shoeboxes)
        new_shoeboxes = reflection_table["shoebox"]
        for s in range(len(new_shoeboxes)):
            bbox_old = strong_shoeboxes[s].bbox
            mask_old = strong_shoeboxes[s].mask
            bbox_new = new_shoeboxes[s].bbox
            mask_new = new_shoeboxes[s].mask
            for j in range(mask_old.all()[1]):
                for i in range(mask_old.all()[2]):
                    ii = bbox_old[0] + i - bbox_new[0]
                    jj = bbox_old[2] + j - bbox_new[2]
                    if mask_old[0, j, i] == 5:
                        if (
                            ii >= 0
                            and jj >= 0
                            and jj < mask_new.all()[1]
                            and ii < mask_new.all()[2]
                        ):
                            assert mask_new[0, jj, ii] & (1 << 0)
                            mask_new[0, jj, ii] |= (1 << 2) | (1 << 3)


def initial_integrator(experiments, reflection_table):
    """Performs an initial integration of strong spots"""

    # some functions require an experimentlist, others just the experiment
    experiment = experiments[0]
    sel = reflection_table.get_flags(reflection_table.flags.strong)
    strong_refls = reflection_table.select(sel)
    strong_shoeboxes = copy.deepcopy(
        strong_refls["shoebox"]
    )  # Save the strong shoeboxes

    # Compute and initial spot size estimate and beam vector
    sigma_d = ComputeEsdBeamDivergence(experiment.detector, strong_refls).sigma()
    sigma_degrees = sigma_d * 180 / pi
    logger.info(
        f"Initial sigma d estimate for {len(strong_refls)} reflections\n"
        + f"Sigma D: {sigma_degrees:.5f} degrees\n",
    )
    strong_refls["s1_obs"] = _compute_beam_vector(experiment, strong_refls)
    strong_refls["bbox"] = _compute_bbox(experiment, strong_refls, sigma_d, "s1_obs")

    # allocate and extract shoebox
    strong_refls["shoebox"] = flex.shoebox(
        strong_refls["panel"], strong_refls["bbox"], allocate=True
    )
    strong_refls.extract_shoeboxes(experiment.imageset)
    reflection_table.is_overloaded(experiments)
    reflection_table.contains_invalid_pixels()

    _compute_mask(experiment, strong_refls, sigma_d, "s1_obs", strong_shoeboxes)

    logger.info(
        f"Computing background, intensity and centroids for {len(strong_refls)} reflections"
    )
    strong_refls.compute_background(experiments)
    strong_refls.compute_summed_intensity()

    n_sum = strong_refls.get_flags(strong_refls.flags.integrated_sum).count(True)
    logger.info(f"{n_sum} reflections integrated")

    strong_refls.compute_centroid(experiments)
    strong_refls.compute_d(experiments)

    return strong_refls, sigma_d


def final_integrator(
    experiments,
    reflection_table,
    sigma_d,
    use_crude_shoebox_mask=False,
    shoebox_probability=FULL_PARTIALITY,
):
    """Performs an initial integration of all predicted spots"""

    experiment = experiments[0]
    logger.info("\n" + "=" * 80 + "\nIntegrating reflections")

    # first compute the bbox
    if use_crude_shoebox_mask:
        reflection_table["bbox"] = _compute_bbox(
            experiment, reflection_table, sigma_d, "s1"
        )
    else:
        # compute bbox from model
        profile = experiment.crystal.mosaicity
        profile.parameterisation.compute_bbox(
            experiments, reflection_table, shoebox_probability
        )

    # Select reflections within detector
    x0, x1, y0, y1, _, _ = reflection_table["bbox"].parts()
    sel = flex.bool(reflection_table.size(), False)
    panel_id = reflection_table["panel"]
    for i in range(len(experiment.detector)):
        panel_sel = panel_id == i
        xsize, ysize = experiment.detector[i].get_image_size()
        selection = panel_sel & (x1 > 0) & (y1 > 0) & (x0 < xsize) & (y0 < ysize)
        sel |= selection
    reflection_table = reflection_table.select(sel)

    reflection_table["shoebox"] = flex.shoebox(
        reflection_table["panel"], reflection_table["bbox"], allocate=True
    )
    reflection_table.extract_shoeboxes(experiment.imageset)

    reflection_table.is_overloaded(experiments)
    reflection_table.contains_invalid_pixels()

    if use_crude_shoebox_mask:
        _compute_mask(experiment, reflection_table, sigma_d, "s1")
    else:
        # compute the mask from the model.
        profile = experiment.crystal.mosaicity
        profile.parameterisation.compute_mask(
            experiments, reflection_table, shoebox_probability
        )

    logger.info(
        f"Computing background, intensity, corrections for {len(reflection_table)} reflections"
    )
    reflection_table.compute_background(experiments)
    reflection_table.compute_summed_intensity()
    reflection_table.compute_corrections(experiments)
    reflection_table.compute_centroid(experiments)

    profile = experiment.crystal.mosaicity
    profile.parameterisation.compute_partiality(experiments, reflection_table)

    return reflection_table


def refine_profile(
    experiment,
    profile,
    refiner_data,
    wavelength_spread_model="delta",
    max_iter=1000,
    LL_tolerance=1e-6,
):
    """Do the profile refinement"""
    logger.info("\n" + "=" * 80 + "\nRefining profile parameters")

    # Create the parameterisation
    state = ModelState(
        experiment,
        profile.parameterisation.parameterisation(),
        fix_orientation=True,
        fix_unit_cell=True,
        fix_wavelength_spread=wavelength_spread_model == "delta",
    )

    # Create the refiner and refine
    refiner = ProfileRefiner(state, refiner_data, max_iter, LL_tolerance)
    refiner.refine()

    # Set the profile parameters
    profile.parameterisation.update_model(refiner.state)
    # Set the mosaicity
    experiment.crystal.mosaicity = profile

    return refiner


def refine_crystal(
    experiment,
    profile,
    refiner_data,
    fix_unit_cell=False,
    fix_orientation=False,
    wavelength_spread_model="delta",
    max_iter=1009,
    LL_tolerance=1e-6,
    max_cell_volume_change_fraction=0.2,
):
    """Do the crystal refinement"""
    if (fix_unit_cell is True) and (fix_orientation is True):
        return

    initial_cell_volume = experiment.crystal.get_unit_cell().volume()
    logger.info("\n" + "=" * 80 + "\nRefining crystal parameters")

    # Create the parameterisation
    state = ModelState(
        experiment,
        profile.parameterisation.parameterisation(),
        fix_mosaic_spread=True,
        fix_unit_cell=fix_unit_cell,
        fix_orientation=fix_orientation,
        fix_wavelength_spread=wavelength_spread_model == "delta",
    )

    # Create the refiner and refine
    refiner = ProfileRefiner(state, refiner_data, max_iter, LL_tolerance)
    refiner.refine()

    end_cell_volume = refiner.state.crystal.get_unit_cell().volume()
    volume_change_fraction = (
        abs(end_cell_volume - initial_cell_volume) / initial_cell_volume
    )
    if volume_change_fraction > max_cell_volume_change_fraction:
        raise RuntimeError(
            f"Cell volume change fraction ({volume_change_fraction:.2}) greater than {max_cell_volume_change_fraction} (max_cell_volume_change_fraction) in this cycle"
            + f"\n  Initial cell volume {initial_cell_volume:.2f}, final cell volume {end_cell_volume:.2f}"
        )

    return refiner


def predict_after_ellipsoid_refinement(experiment, reflection_table):
    """
    Predict the position of the spots

    """

    # Compute the vector to the reciprocal lattice point
    # since this is not on the ewald sphere, lets call it s2
    s1, s2 = calc_s1_s2(
        reflection_table["miller_index"],
        experiment.crystal.get_A(),
        experiment.beam.get_s0(),
    )
    reflection_table["s1"] = s1
    reflection_table["s2"] = s2
    reflection_table["entering"] = flex.bool(reflection_table.size(), False)

    # Compute the ray intersections
    xyzpx = flex.vec3_double()
    xyzmm = flex.vec3_double()
    for ss, panel_id in zip(s1, reflection_table["panel"]):
        mm = experiment.detector[panel_id].get_ray_intersection(ss)
        px = experiment.detector[panel_id].millimeter_to_pixel(mm)
        xyzpx.append(px + (0,))
        xyzmm.append(mm + (0,))
    reflection_table["xyzcal.mm"] = xyzmm
    reflection_table["xyzcal.px"] = xyzpx
    return reflection_table


def compute_prediction_probability(experiment, reflection_table):
    # Get stuff from experiment
    s0 = np.array([experiment.beam.get_s0()], dtype=np.float64).reshape(3, 1)
    s0_length = norm(s0)
    profile = experiment.crystal.mosaicity

    # Loop through reflections
    min_p = None
    for s2 in reflection_table["s2"]:
        s2 = np.array(list(s2), dtype=np.float64).reshape(3, 1)
        s3 = s2 * s0_length / norm(s2)
        r = s2 - s0
        epsilon = s3 - s2
        sigma = profile.parameterisation.sigma_for_reflection(s0, r)
        sigma_inv = inv(sigma)
        d = np.matmul(np.matmul(epsilon.T, sigma_inv), epsilon)[0, 0]
        p = chisq_pdf(3, float(d))
        if min_p is None or p < min_p:
            min_p = p

    # Print some stuff
    logger.info(
        "Quantile required to predicted all observed reflections = %.5f" % (1 - min_p)
    )


def run_ellipsoid_refinement(
    experiments,
    reflection_table,
    sigma_d,
    profile_model="simple6",
    wavelength_model="delta",
    fix_unit_cell=False,
    fix_orientation=False,
    capture_progress=False,
    n_cycles=3,
    max_iter=1000,
    LL_tolerance=1e-6,
    mosaicity_max_limit=0.004,
    max_cell_volume_change_fraction=0.2,
):
    """Runs ellipsoid refinement on strong spots.

    Creates the necessary data needed, then runs cycles of profile and crystal
    refinement,"""

    output_data = {
        "refiner_output": {
            "history": [],
            "correlation": None,
            "labels": None,
        },
    }

    # Set the M params
    if not hasattr(experiments[0].crystal, "mosaicity"):
        profile = ProfileModelFactory.from_sigma_d(profile_model, sigma_d)
    else:
        profile = experiments[0].crystal.mosaicity
    profile.parameterisation.mosaicity_max_limit = mosaicity_max_limit
    # Construct the profile refiner data
    refiner_data = RefinerData.from_reflections(experiments[0], reflection_table)
    profile.parameterisation.n_obs = len(refiner_data.h_list)
    # Do the refinement
    for _ in range(n_cycles):
        # refine the profile
        refiner = refine_profile(
            experiments[0],
            profile,
            refiner_data,
            wavelength_spread_model=wavelength_model,
            max_iter=max_iter,
            LL_tolerance=LL_tolerance,
        )
        if capture_progress:
            # Save some data for plotting later.
            output_data["refiner_output"]["history"].append(refiner.history)
            output_data["refiner_output"]["correlation"] = refiner.correlation()
            output_data["refiner_output"]["labels"] = refiner.labels()

        # refine the crystal
        refinerc = refine_crystal(
            experiments[0],
            profile,
            refiner_data,
            fix_unit_cell=fix_unit_cell,
            fix_orientation=fix_orientation,
            wavelength_spread_model=wavelength_model,
            max_iter=max_iter,
            LL_tolerance=LL_tolerance,
            max_cell_volume_change_fraction=max_cell_volume_change_fraction,
        )
        if capture_progress:
            # Save some data for plotting later.
            output_data["refiner_output"]["history"].append(refinerc.history)
            output_data["refiner_output"]["correlation"] = refinerc.correlation()
            output_data["refiner_output"]["labels"] = refinerc.labels()

    experiments[0].profile = profile

    # Post process the reflections
    # Update predictions
    reflection_table = predict_after_ellipsoid_refinement(
        experiments[0], reflection_table
    )
    # Compute prob
    compute_prediction_probability(experiments[0], reflection_table)

    return experiments, reflection_table, output_data


def predict(experiments, d_min=None, prediction_probability=FULL_PARTIALITY):
    """Predict the reflections"""
    logger.info("\n" + "=" * 80 + "\nPredicting reflections")

    # Set a resolution range
    if d_min is None:
        s0 = experiments[0].beam.get_s0()
        d_min = experiments[0].detector.get_max_resolution(s0)

    # Create the index generator
    index_generator = IndexGenerator(
        experiments[0].crystal.get_unit_cell(),
        experiments[0].crystal.get_space_group().type(),
        d_min,
    )

    # Get an array of miller indices
    miller_indices_to_test = index_generator.to_array()
    logger.info("Generated %d miller indices" % len(miller_indices_to_test))

    # Get the covariance matrix
    profile = experiments[0].crystal.mosaicity

    reflection_table = profile.parameterisation.predict_reflections(
        experiments, miller_indices_to_test, prediction_probability
    )

    # Do the prediction
    reflection_table.compute_d(experiments)
    logger.info("Predicted %d reflections" % len(reflection_table))

    return reflection_table
