#!/usr/bin/env python
"""
This script optimizes Miller indices and wavelengths jointly.
"""

import logging
import sys
import time
from itertools import repeat
from multiprocessing import Pool

import libtbx.phil
import numpy as np
from dials.array_family.flex import reflection_table
from dials.util import show_mail_handle_errors
from dials.util.options import (ArgumentParser,
                                reflections_and_experiments_from_files)
from dxtbx.model import ExperimentList

logger = logging.getLogger("laue-dials.command_line.optimize_indexing")

help_message = """

This program takes initial monochromatic estimates of the geometry in a
DIALS experiment list and reflection table, and optimizes the indexed
solution by allowing the wavelength to vary. Only Miller indices, s1
vectors and wavelengths for scattered reflections are overwritten by
this program - all other geometry variables remain constant here.

The outputs are a pair of files (optimized.expt, optimized.refl), which
mimic the input files but with updated Miller indices, s1 vectors, and
wavelengths in the reflection table. Note that the experiment list is
unchanged entirely.

Examples::

    laue.optimize_indexing [options] monochromatic.expt monochromatic.refl
"""

# Set the phil scope
phil_scope = libtbx.phil.parse(
    """

output {
  experiments = 'optimized.expt'
    .type = str
    .help = "The output experiment list filename."

  reflections = 'optimized.refl'
    .type = str
    .help = "The output reflection table filename."

  log = 'laue.optimize_indexing.log'
    .type = str
    .help = "The log filename."
  }

keep_unindexed = False
  .type = bool
  .help = Whether to keep unindexed reflections

n_proc = 1
  .type = int
  .help = Number of parallel processes to run

n_macrocycles = 3
  .type = int(value_min=1)
  .help = "Number of macrocycles of index optimization to perform"

wavelengths {
  lam_min = None
    .type = float(value_min=0.1)
    .help = "Minimum wavelength for beam spectrum"
  lam_max = None
    .type = float(value_min=0.2)
    .help = "Maximum wavelength for beam spectrum"
  }

reciprocal_grid {
  d_min = None
    .type = float(value_min=0.1)
    .help = "Minimum d-spacing for reflecting planes"
  }

""",
    process_includes=True,
)

working_phil = phil_scope.fetch(sources=[phil_scope])


def index_image(params, refls, expts):
    import gemmi
    from cctbx.sgtbx import space_group
    from cctbx.uctbx import unit_cell
    from dials.array_family import flex
    from dxtbx.model import ExperimentList

    from laue_dials.algorithms.laue import LaueAssigner

    if type(expts) != ExperimentList:
        expts = ExperimentList([expts])

    s0 = np.array(expts[0].beam.get_s0())

    # Write to reflection file
    refls["wavelength"] = flex.double(len(refls))
    refls["miller_index"] = flex.miller_index(len(refls))
    refls["harmonics"] = flex.bool([False] * len(refls))

    for i in range(len(expts.imagesets())):
        # Get experiment data from experiment objects
        experiment = expts[i]
        cryst = experiment.crystal
        spacegroup = gemmi.SpaceGroup(
            cryst.get_space_group().type().universal_hermann_mauguin_symbol()
        )

        # Get reflections on this image
        idx = refls["id"] == int(experiment.identifier)
        subrefls = refls.select(idx)

        # Get unit cell params
        cell_params = cryst.get_unit_cell().parameters()
        cell = gemmi.UnitCell(*cell_params)

        # Generate s vectors
        s1 = subrefls["s1"].as_numpy_array()

        # Get U matrix
        U = np.asarray(cryst.get_U()).reshape(3, 3)

        # Generate assigner object
        logger.info(f"Reindexing image {experiment.identifier}.")
        la = LaueAssigner(
            s0,
            s1,
            cell,
            U,
            params.wavelengths.lam_min,
            params.wavelengths.lam_max,
            params.reciprocal_grid.d_min,
            spacegroup,
        )

        # Optimize Miller indices
        la.assign()
        for j in range(params.n_macrocycles):
            la.reset_inliers()
            la.update_rotation()
            la.assign()
            la.reject_outliers()
            la.update_rotation()
            la.assign()

        # Update s1 based on new wavelengths
        s1[la._inliers] = la.s1 / la.wav[:, None]

        # Reset crystal parameters based on new geometry
        cryst.set_U(la.R.flatten())
        cryst.set_A(la.RB.flatten())
        cryst.set_B(la.B.flatten())
        cryst.set_space_group(space_group(la.spacegroup.hall))
        cryst.set_unit_cell(unit_cell(la.cell.parameters))

        # Filter out wavelengths beyond limits
        spot_wavelengths = np.asarray(la._wav.tolist())
        spot_wavelengths[spot_wavelengths < params.wavelengths.lam_min] = 0
        spot_wavelengths[spot_wavelengths > params.wavelengths.lam_max] = 0

        # Write data to reflections
        refls["s1"].set_selected(idx, flex.vec3_double(s1))
        refls["miller_index"].set_selected(
            idx,
            flex.miller_index(la._H.astype("int").tolist()),
        )
        refls["harmonics"].set_selected(
            idx,
            flex.bool(la._harmonics.tolist()),
        )
        refls["wavelength"].set_selected(
            idx,
            flex.double(spot_wavelengths),
        )

    # Remove unindexed reflections
    if not params.keep_unindexed:
        all_wavelengths = refls["wavelength"].as_numpy_array()
        keep = np.logical_and(
            all_wavelengths >= params.wavelengths.lam_min,
            all_wavelengths <= params.wavelengths.lam_max,
        )
        refls = refls.select(flex.bool(keep))

    # Return reindexed expts, refls
    return expts, refls


@show_mail_handle_errors()
def run(args=None, *, phil=working_phil):
    # Parse arguments
    usage = "laue.optimize_indexing [options] monochromatic.expt monochromatic.refl"

    parser = ArgumentParser(
        usage=usage,
        phil=phil,
        read_reflections=True,
        read_experiments=True,
        check_format=False,
        epilog=help_message,
    )

    params, options = parser.parse_args(args=args, show_diff_phil=False)

    # Configure logging
    console = logging.StreamHandler(sys.stdout)
    fh = logging.FileHandler(params.output.log, mode="w", encoding="utf-8")
    loglevel = logging.INFO

    logger.addHandler(fh)
    logger.addHandler(console)
    logging.captureWarnings(True)
    warning_logger = logging.getLogger("py.warnings")
    warning_logger.addHandler(fh)
    warning_logger.addHandler(console)
    dials_logger = logging.getLogger("dials")
    dials_logger.addHandler(fh)
    dials_logger.addHandler(console)
    dxtbx_logger = logging.getLogger("dxtbx")
    dxtbx_logger.addHandler(fh)
    dxtbx_logger.addHandler(console)
    xfel_logger = logging.getLogger("xfel")
    xfel_logger.addHandler(fh)
    xfel_logger.addHandler(console)

    logger.setLevel(loglevel)
    dials_logger.setLevel(loglevel)
    dxtbx_logger.setLevel(loglevel)
    xfel_logger.setLevel(loglevel)
    fh.setLevel(loglevel)

    # Log diff phil
    diff_phil = parser.diff_phil.as_str()
    if diff_phil != "":
        logger.info("The following parameters have been modified:\n")
        logger.info(diff_phil)

    # Print help if no input
    if not params.input.experiments or not params.input.reflections:
        parser.print_help()
        return

    reflections, experiments = reflections_and_experiments_from_files(
        params.input.reflections, params.input.experiments
    )
    reflections = reflections[0]  # Get reflection table out of list

    # Sanity checks
    if len(experiments) == 0:
        parser.print_help()
        return

    # Check for valid parameter values
    if params.reciprocal_grid.d_min == None:
        logger.info("Please provide a d_min.")
        return
    elif params.wavelengths.lam_min == None or params.wavelengths.lam_max == None:
        logger.info(
            "Please provide upper and lower boundaries for the wavelength spectrum."
        )
        return
    elif params.wavelengths.lam_min > params.wavelengths.lam_max:
        logger.info("Minimum wavelength cannot be greater than maximum wavelength.")
        return

    # Get initial time for process
    start_time = time.time()

    # This will populate ['s1'] & ['rlp'] columns
    reflections.centroid_px_to_mm(experiments)
    reflections.map_centroids_to_reciprocal_space(experiments)

    # Prepare parallel input
    ids = list(np.unique(reflections["id"]).astype(np.int32))
    expts_arr = []
    refls_arr = []
    for i in ids:  # Split DIALS objects into lists
        expts_arr.append(ExperimentList([experiments[i]]))
        refls_arr.append(reflections.select(reflections["id"] == i))
    inputs = list(zip(repeat(params), refls_arr, expts_arr))

    # Reindex data
    num_processes = params.n_proc
    logger.info("Reindexing images.")
    with Pool(processes=num_processes) as pool:
        output = pool.starmap(index_image, inputs, chunksize=1)
    logger.info(f"All images reindexed.")

    # Convert reindexed data to DIALS objects
    total_experiments = ExperimentList()
    total_reflections = reflection_table()
    for i in ids:
        total_experiments.extend(output[i][0])
        total_reflections.extend(output[i][1])

    # Save experiments
    logger.info("Saving optimized experiments to %s", params.output.experiments)
    total_experiments.as_file(params.output.experiments)

    # Save reflections
    logger.info("Saving optimized reflections to %s", params.output.reflections)
    total_reflections.as_file(filename=params.output.reflections)

    # Final logs
    logger.info("")
    logger.info(
        "Time Taken for Total Processing = %f seconds", time.time() - start_time
    )


if __name__ == "__main__":
    run()
