// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ImgCIF.h"

#include <algorithm>
#include <cctype>
#include <cmath>
#include <set>

#include "../common/JFJochException.h"
#include "../common/DiffractionGeometry.h"
#include "../common/JFJochMath.h"

namespace imgcif {

namespace {

using Vec = std::array<double, 3>;

JFJochException Invalid(const std::string &what) {
    return {JFJochExceptionCategory::InputParameterInvalid, "imgCIF header: " + what};
}

std::string Field(const std::map<std::string, std::string> &row, const std::string &name) {
    const auto it = row.find(name);
    return (it == row.end()) ? std::string() : it->second;
}

// A CIF number. "?" (unknown) and "." (not applicable) are how a table leaves a cell empty - an angle
// column on a translation row - and read as nothing.
std::optional<double> Number(const std::string &s) {
    if (s.empty() || s == "?" || s == ".")
        return {};
    try {
        size_t used = 0;
        const double v = std::stod(s, &used);
        // CIF allows a standard uncertainty in brackets after the value: "0.71073(2)"
        if (used != s.size() && s[used] != '(')
            throw Invalid("malformed number '" + s + "'");
        return v;
    } catch (const std::logic_error &) {
        throw Invalid("malformed number '" + s + "'");
    }
}

double NumberOr(const std::string &s, double fallback) {
    return Number(s).value_or(fallback);
}

Vec Rotate(double angle_deg, Vec axis, const Vec &p) {
    const double len = std::sqrt(axis[0] * axis[0] + axis[1] * axis[1] + axis[2] * axis[2]);
    for (double &c: axis)
        c /= len;
    const double t = angle_deg * 3.14159265358979323846 / 180.0;
    const double c = std::cos(t), s = std::sin(t);
    const double dot = axis[0] * p[0] + axis[1] * p[1] + axis[2] * p[2];
    const Vec cross = {axis[1] * p[2] - axis[2] * p[1],
                       axis[2] * p[0] - axis[0] * p[2],
                       axis[0] * p[1] - axis[1] * p[0]};
    Vec out{};
    for (int i = 0; i < 3; i++)
        out[i] = p[i] * c + cross[i] * s + axis[i] * dot * (1.0 - c);
    return out;
}

const Axis &Find(const Description &d, const std::string &id) {
    const auto it = d.axes.find(id);
    if (it == d.axes.end())
        throw Invalid("axis '" + id + "' is used but not defined in the _axis table");
    return it->second;
}

// Carry a point from the frame of axis `id` out to the laboratory: each axis on the way applies its own
// motion at its setting (or at the one given in `settings`), then moves by its offset from the axis it
// depends on. A child's vector is stated in its parent's frame, which is exactly what applying the
// parent after the child does to it.
Vec ToLab(const Description &d, std::string id, Vec p, const std::map<std::string, double> &settings) {
    std::set<std::string> seen;
    while (id != ".") {
        if (!seen.insert(id).second)
            throw Invalid("the axis chain through '" + id + "' loops back on itself");
        const Axis &a = Find(d, id);
        const auto s = settings.find(id);
        const double value = (s != settings.end()) ? s->second : a.setting;
        if (a.type == "rotation")
            p = Rotate(value, a.vector, p);
        else if (a.type == "translation")
            for (int i = 0; i < 3; i++)
                p[i] += value * a.vector[i];
        else
            throw Invalid("axis '" + id + "' is of type '" + a.type + "', which is not supported");
        for (int i = 0; i < 3; i++)
            p[i] += a.offset[i];
        id = a.depends_on;
    }
    return p;
}

// Where stored pixel `stored` (0-based) sits along its axis. _array_structure_list.direction is not
// consulted: the dictionary says "decreasing" means the stored order runs from the last index to the
// first, but CBFlib - and every program built on it - places a pixel by the axis displacement alone,
// and the files written through CBFlib are only consistent read that way. Bruker's conversions state
// "decreasing" for their rows: read as CBFlib reads them a sweep of them indexes (and DIALS indexes
// it), read by the dictionary the image is mirrored and nothing does.
double Displacement(const ArrayIndex &index, double stored) {
    return index.displacement + stored * index.increment;
}

}  // namespace

std::vector<std::map<std::string, std::string>> ParseLoop(const std::string &text, const std::string &tag) {
    // The loop_ that introduces the tag, not the tag's own position: everything before it is another
    // loop's data.
    const size_t tag_at = text.find("\n" + tag);
    if (tag_at == std::string::npos)
        return {};
    const size_t loop_at = text.rfind("loop_", tag_at);
    if (loop_at == std::string::npos)
        return {};

    // Bare tokens, and the quoted ones a CIF value may be - a quoted value holding spaces would
    // otherwise be counted as several columns and shift every row after it.
    std::vector<std::string> tokens;
    for (size_t i = loop_at + 5; i < text.size();) {
        while ((i < text.size()) && std::isspace(static_cast<unsigned char>(text[i])))
            i++;
        if (i >= text.size())
            break;
        size_t end;
        if ((text[i] == '\'') || (text[i] == '"')) {
            end = text.find(text[i], i + 1);
            if (end == std::string::npos)
                break;
            tokens.push_back(text.substr(i + 1, end - i - 1));
            end++;
        } else {
            end = i;
            while ((end < text.size()) && !std::isspace(static_cast<unsigned char>(text[end])))
                end++;
            tokens.push_back(text.substr(i, end - i));
        }
        // A second loop_, or a tag belonging to another category, ends this one.
        if ((tokens.back() == "loop_")
                || (tokens.back().starts_with("_") && !tokens.back().starts_with(tag.substr(0, tag.find('.') + 1)))) {
            tokens.pop_back();
            break;
        }
        i = end;
    }

    std::vector<std::string> names;
    size_t first_value = 0;
    while ((first_value < tokens.size()) && tokens[first_value].starts_with("_"))
        names.push_back(tokens[first_value++]);
    // The loop found must be the tag's own: a tag stated outside any loop would otherwise be read as
    // the rows of whatever loop came before it.
    if (std::find(names.begin(), names.end(), tag) == names.end())
        return {};

    std::vector<std::map<std::string, std::string>> rows;
    for (size_t i = first_value; i + names.size() <= tokens.size(); i += names.size()) {
        std::map<std::string, std::string> row;
        for (size_t j = 0; j < names.size(); j++)
            row[names[j]] = tokens[i + j];
        rows.push_back(std::move(row));
    }
    return rows;
}

std::optional<std::string> Item(const std::string &text, const std::string &tag) {
    const std::string key = "\n" + tag;
    for (size_t at = text.find(key); at != std::string::npos; at = text.find(key, at + 1)) {
        size_t i = at + key.size();
        // the whole tag, not a longer one that starts with it
        if ((i < text.size()) && !std::isspace(static_cast<unsigned char>(text[i])))
            continue;
        while ((i < text.size()) && std::isspace(static_cast<unsigned char>(text[i])))
            i++;
        if (i >= text.size())
            return {};
        if ((text[i] == '\'') || (text[i] == '"')) {
            const size_t end = text.find(text[i], i + 1);
            if (end == std::string::npos)
                return {};
            return text.substr(i + 1, end - i - 1);
        }
        size_t end = i;
        while ((end < text.size()) && !std::isspace(static_cast<unsigned char>(text[end])))
            end++;
        const std::string value = text.substr(i, end - i);
        if (value.starts_with("_") || value == "loop_")
            return {};          // the tag heads a loop: not a single item
        return value;
    }
    return {};
}

std::optional<Description> Parse(const std::string &text) {
    const auto axis_rows = ParseLoop(text, "_axis.id");
    auto frame_rows = ParseLoop(text, "_diffrn_scan_frame_axis.axis_id");
    const auto scan_rows = ParseLoop(text, "_diffrn_scan_axis.axis_id");
    if (axis_rows.empty() || (frame_rows.empty() && scan_rows.empty()))
        return {};

    Description d;
    for (const auto &row: axis_rows) {
        Axis a;
        a.type = Field(row, "_axis.type");
        a.equipment = Field(row, "_axis.equipment");
        a.depends_on = Field(row, "_axis.depends_on");
        if (a.depends_on.empty() || a.depends_on == "?")
            a.depends_on = ".";
        for (int i = 0; i < 3; i++) {
            a.vector[i] = NumberOr(Field(row, "_axis.vector[" + std::to_string(i + 1) + "]"), 0.0);
            a.offset[i] = NumberOr(Field(row, "_axis.offset[" + std::to_string(i + 1) + "]"), 0.0);
        }
        d.axes[Field(row, "_axis.id")] = a;
    }

    // Where each axis stands: the frame's own setting where the file states one, the scan's start
    // otherwise; how far it moves over the frame, from the scan.
    for (const auto &row: scan_rows) {
        const auto it = d.axes.find(Field(row, "_diffrn_scan_axis.axis_id"));
        if (it == d.axes.end())
            continue;
        const bool rotation = it->second.type == "rotation";
        it->second.setting = NumberOr(Field(row, rotation ? "_diffrn_scan_axis.angle_start"
                                                          : "_diffrn_scan_axis.displacement_start"), 0.0);
        it->second.increment = NumberOr(Field(row, rotation ? "_diffrn_scan_axis.angle_increment"
                                                            : "_diffrn_scan_axis.displacement_increment"), 0.0);
    }
    for (const auto &row: frame_rows) {
        const auto it = d.axes.find(Field(row, "_diffrn_scan_frame_axis.axis_id"));
        if (it == d.axes.end())
            continue;
        const bool rotation = it->second.type == "rotation";
        if (const auto v = Number(Field(row, rotation ? "_diffrn_scan_frame_axis.angle"
                                                      : "_diffrn_scan_frame_axis.displacement")))
            it->second.setting = *v;
    }

    // The two image indices, and the axis each runs along.
    std::map<std::string, std::map<std::string, std::string>> set_rows;
    for (const auto &row: ParseLoop(text, "_array_structure_list_axis.axis_set_id"))
        set_rows[Field(row, "_array_structure_list_axis.axis_set_id")] = row;
    for (const auto &row: ParseLoop(text, "_array_structure_list.array_id")) {
        ArrayIndex index;
        const std::string set = Field(row, "_array_structure_list.axis_set_id");
        const auto s = set_rows.find(set);
        if (s == set_rows.end())
            throw Invalid("image axis set '" + set + "' has no _array_structure_list_axis row");
        index.axis_id = Field(s->second, "_array_structure_list_axis.axis_id");
        index.displacement = NumberOr(Field(s->second, "_array_structure_list_axis.displacement"), 0.0);
        index.increment = NumberOr(Field(s->second, "_array_structure_list_axis.displacement_increment"), 0.0);
        index.dimension = static_cast<int64_t>(NumberOr(Field(row, "_array_structure_list.dimension"), 0.0));
        const std::string which = Field(row, "_array_structure_list.index");
        if (which == "1")
            d.fast = index;
        else if (which == "2")
            d.slow = index;
    }
    if (d.fast.axis_id.empty() || d.slow.axis_id.empty())
        throw Invalid("no _array_structure_list rows for the two image indices");

    d.wavelength_A = NumberOr(Item(text, "_diffrn_radiation_wavelength.wavelength").value_or(""), 0.0);
    if (d.wavelength_A == 0.0) {
        const auto rows = ParseLoop(text, "_diffrn_radiation_wavelength.id");
        if (!rows.empty())
            d.wavelength_A = NumberOr(Field(rows.front(), "_diffrn_radiation_wavelength.wavelength"), 0.0);
    }
    d.overload = NumberOr(Item(text, "_array_intensities.overload").value_or(""), 0.0);
    d.exposure_s = NumberOr(Item(text, "_diffrn_scan_frame.integration_time").value_or(""), 0.0);
    d.detector = Item(text, "_diffrn_detector.type").value_or(Item(text, "_diffrn_detector.detector").value_or(""));
    d.ccd = Item(text, "_diffrn_detector.detector").value_or("").find("CCD") != std::string::npos;
    if (const auto ratio = Item(text, "_diffrn_radiation_wavelength.polarizn_source_ratio"))
        d.polarization_ratio = NumberOr(*ratio, 0.0);
    d.polarization_norm_deg = NumberOr(Item(text, "_diffrn_radiation_wavelength.polarizn_source_norm").value_or(""), 0.0);
    return d;
}

std::optional<float> PolarizationFactor(const Description &d) {
    if (!d.polarization_ratio)
        return {};
    return static_cast<float>(*d.polarization_ratio * std::cos(2.0 * d.polarization_norm_deg * PI / 180.0));
}

Geometry Evaluate(const Description &d) {
    Geometry g;

    // The image: the index axis that hangs off the other one is where the chain starts.
    const Axis &fast = Find(d, d.fast.axis_id);
    const Axis &slow = Find(d, d.slow.axis_id);
    const bool slow_is_leaf = slow.depends_on == d.fast.axis_id;
    if (!slow_is_leaf && fast.depends_on != d.slow.axis_id)
        throw Invalid("the two image axes " + d.fast.axis_id + " and " + d.slow.axis_id
                      + " do not depend on each other; such a detector is not supported");
    const std::string leaf = slow_is_leaf ? d.slow.axis_id : d.fast.axis_id;
    const auto pixel = [&](double col, double row) {
        return ToLab(d, leaf, {0, 0, 0}, {{d.fast.axis_id, Displacement(d.fast, col)},
                                          {d.slow.axis_id, Displacement(d.slow, row)}});
    };
    g.origin = pixel(0, 0);
    const Vec p10 = pixel(1, 0), p01 = pixel(0, 1);
    for (int i = 0; i < 3; i++) {
        g.fast[i] = p10[i] - g.origin[i];
        g.slow[i] = p01[i] - g.origin[i];
    }

    // The scan: exactly one axis may move, and it has to be a goniometer rotation. Its direction is
    // its own vector carried by every axis it depends on, at their settings - on a kappa head a phi
    // scan turns about phi as omega and kappa hold it. The axes it carries do not move it.
    std::vector<std::string> moving;
    for (const auto &[id, a]: d.axes)
        if (a.increment != 0.0 && id != d.fast.axis_id && id != d.slow.axis_id)
            moving.push_back(id);
    if (moving.size() > 1) {
        std::string names;
        for (const auto &m: moving)
            names += (names.empty() ? "" : ", ") + m;
        throw Invalid("more than one axis moves during the frame (" + names + "); only a rotation "
                      "about one goniometer axis is supported");
    }
    if (moving.empty())
        return g;

    const std::string &id = moving.front();
    const Axis &a = d.axes.at(id);
    if (a.type != "rotation" || a.equipment != "goniometer")
        throw Invalid("the axis that moves during the frame (" + id + ") is a " + a.equipment + " "
                      + a.type + "; only a goniometer rotation is supported");

    // Rotations only: an offset moves where the axis is, not which way it points.
    Vec dir = a.vector;
    std::set<std::string> seen;
    for (std::string up = a.depends_on; up != "."; up = Find(d, up).depends_on) {
        if (!seen.insert(up).second)
            throw Invalid("the axis chain through '" + up + "' loops back on itself");
        const Axis &parent = Find(d, up);
        if (parent.type == "rotation")
            dir = Rotate(parent.setting, parent.vector, dir);
    }
    const double len = std::sqrt(dir[0] * dir[0] + dir[1] * dir[1] + dir[2] * dir[2]);
    if (!(len > 0.0))
        throw Invalid("the scan axis " + id + " has no direction");
    for (double &c: dir)
        c /= len;

    g.scan_axis = id;
    g.scan_direction = dir;
    g.scan_start_deg = a.setting;
    g.scan_increment_deg = a.increment;
    return g;
}

Coord ToInternal(const std::array<double, 3> &v) {
    return {static_cast<float>(v[0]), static_cast<float>(-v[1]), static_cast<float>(-v[2])};
}

Coord GoniometerAxisOf(const std::array<double, 3> &v) {
    return -ToInternal(v);
}

GoniometerAxis SweepAxis(const std::string &name, double start_deg, double increment_deg,
                         const std::array<double, 3> &v) {
    const double sign = increment_deg < 0.0 ? -1.0 : 1.0;
    return GoniometerAxis(name, static_cast<float>(sign * start_deg), static_cast<float>(sign * increment_deg),
                          GoniometerAxisOf(v) * static_cast<float>(sign), {});
}

DetectorPlane PlaneOf(const Geometry &g, const std::string &path) {
    const Coord origin = ToInternal(g.origin);
    const Coord fast = ToInternal(g.fast);
    const Coord slow = ToInternal(g.slow);

    DetectorPlane p;
    p.pixel_mm = fast.Length();
    // One pixel size and a rectangular grid is what the geometry can describe.
    if (std::fabs(slow.Length() - p.pixel_mm) > 1e-4f * p.pixel_mm
            || std::fabs(fast * slow) > 1e-4f * p.pixel_mm * p.pixel_mm)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": the imgCIF image axes are not square pixels on a rectangular "
                                     "grid, which is not supported");

    const Coord f = fast.Normalize(), s = slow.Normalize();
    Coord n = f % s;
    // fast x slow points away from the sample on an image stored the internal way round, and towards
    // it on a mirrored one.
    const bool mirrored = n * origin < 0.0f;
    if (mirrored)
        n = -n;

    p.distance_mm = n * origin;
    const Coord to_poni = n * p.distance_mm - origin;
    p.beam_x_pxl = (to_poni * f) / p.pixel_mm;
    p.beam_y_pxl = (to_poni * s) / p.pixel_mm;

    // Of the four ways the image can be stored with this handedness, the one that leaves the least
    // for the continuous rotation - the detector as it is mounted, with the arm and the tilts on top.
    const RotMatrix m(f, s, n);
    float best_trace = -10.0f;
    for (int64_t q = 0; q < 4; q++) {
        const DetectorOrientation o(mirrored, q);
        const RotMatrix r = m * o.Matrix().transpose();
        const float trace = r.Column(0).x + r.Column(1).y + r.Column(2).z;
        if (trace > best_trace) {
            best_trace = trace;
            p.orientation = o;
            PoniAnglesFromMatrix(r, p.rot1, p.rot2, p.rot3);
        }
    }
    return p;
}

}  // namespace imgcif