#!/usr/bin/env python3

import math
import numpy
import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import rcParams
from gi.repository import GLib
import gi

gi.require_version("Hkl", "5.0")
from gi.repository import Hkl


def compute_hkl_trajectories(engines, engine, hkl1=None, hkl2=None, n=100):
    """
    compute all the trajectories for a given engine already configured
    """
    if not hkl1:
        hkl1 = [0, 0, 1]
    if not hkl2:
        hkl2 = [0, 1, 1]

    h = numpy.linspace(hkl1[0], hkl2[0], n + 1)
    k = numpy.linspace(hkl1[1], hkl2[1], n + 1)
    l = numpy.linspace(hkl1[2], hkl2[2], n + 1)

    # set the hkl engine and get the results
    trajectories = []
    for hh, kk, ll in zip(h, k, l):
        try:
            solutions = engine.pseudo_axis_values_set([hh, kk, ll], Hkl.UnitEnum.USER)
            first_solution = solutions.items()[0]
            for i, item in enumerate(solutions.items()):
                try:
                    trajectories[i]
                except IndexError:
                    trajectories.append([])
                values = item.geometry_get().axis_values_get(Hkl.UnitEnum.USER)
                trajectories[i].append(values)
            engines.select_solution(first_solution)
        except GLib.GError as err:
            pass

    return trajectories


def _plot_legend(axes):
    plt.subplot(3, 4, 1)
    plt.title("legend")
    print("legende", 1)
    for name in axes:
        plt.plot([0, 0], label=name)
    plt.legend()


def plot_hkl_trajectory(filename, geometry, engines, hkl1=None, hkl2=None, n=100):
    """
    plot the trajectory for a engine. It is possible to limit the
    number of trajectory using the max_traj keyword
    """
    axis_names = geometry.axis_names_get()

    hkl = engines.engine_get_by_name("hkl")
    page = 1
    plt.clf()
    plt.suptitle(
        '"' + filename + '" ' + repr(hkl1) + " -> " + repr(hkl2) + " page " + str(page)
    )
    _plot_legend(axis_names)
    idx = 2
    for mode in hkl.modes_names_get():
        hkl.current_mode_set(mode)
        trajectories = compute_hkl_trajectories(engines, hkl, hkl1=hkl1, hkl2=hkl2, n=n)
        print('"' + filename + '"', idx, mode, len(trajectories))

        plt.subplot(3, 4, idx)
        plt.title(mode)
        if not len(trajectories):
            plt.text(
                0.5,
                0.5,
                "Failed",
                size=20,
                rotation=0.0,
                ha="center",
                va="center",
                bbox=dict(
                    boxstyle="round",
                    ec=(1.0, 0.5, 0.5),
                    fc=(1.0, 0.8, 0.8),
                ),
            )
            plt.draw()
        else:
            plt.ylim(-180, 180)
            if len(trajectories[0]) == 1:
                plt.plot(trajectories[0], "o-")
            else:
                plt.plot(trajectories[0], "-")

        idx += 1
        if idx > 12:
            pp.savefig()
            plt.clf()
            page += 1
            _plot_legend(axis_names)
            plt.suptitle(
                filename + " " + repr(hkl1) + " -> " + repr(hkl2) + " page " + str(page)
            )
            idx = 2
    pp.savefig()


pp = PdfPages("trajectories.pdf")
rcParams["font.size"] = 6


def main():
    sample = Hkl.Sample.new("toto")
    lattice = Hkl.Lattice.new(
        1.54, 1.54, 1.54, math.radians(90.0), math.radians(90.0), math.radians(90.0)
    )
    sample.lattice_set(lattice)

    detector = Hkl.Detector.factory_new(Hkl.DetectorType(0))

    for key, factory in Hkl.factories().items():
        geometry = factory.create_new_geometry()
        engines = factory.create_new_engine_list()

        # here we set the detector arm with only positiv values for
        # now tth or delta arm
        for axis in geometry.axis_names_get():
            if axis in ["tth", "delta"]:
                tmp = geometry.axis_get(axis)
                tmp.min_max_set(0, 180.0, Hkl.UnitEnum.USER)
                geometry.axis_set(axis, tmp)

        engines.init(geometry, detector, sample)

        engines_names = [engine.name_get() for engine in engines.engines_get()]
        if "hkl" in engines_names:
            plot_hkl_trajectory(
                key, geometry, engines, hkl1=[0, 0, 1], hkl2=[0, 1, 1], n=100
            )
    pp.close()


if __name__ == "__main__":
    main()
