// Copyright (2019-2023) Paul Scherrer Institute

#include <thread>

#include "JFJochStateMachine.h"

void LoadDatasetSettings(DiffractionExperiment& experiment, const DatasetSettings &settings) {
    // Save DatasetSettings - if something goes wrong, restore old settings
    auto tmp = experiment;

    try {
        experiment.ImagesPerTrigger(settings.images_per_trigger);
        experiment.NumTriggers(settings.ntrigger);
        experiment.BeamX_pxl(settings.beam_x_pxl);
        experiment.BeamY_pxl(settings.beam_y_pxl);
        experiment.DetectorDistance_mm(settings.detector_distance_mm);
        experiment.PhotonEnergy_keV(settings.photon_energy_keV);
        experiment.FilePrefix(settings.file_prefix);
        experiment.DataFileCount(settings.data_file_count);
        if (settings.unit_cell)
            experiment.SetUnitCell(settings.unit_cell.value());
        else
            experiment.SetUnitCell();
        experiment.SpaceGroupNumber(settings.space_group_number);
        experiment.SampleName(settings.sample_name);
        experiment.Compression(settings.compression);
        experiment.SaveCalibration(settings.save_calibration);
        if (settings.summation == 0)
            experiment.Summation(1);
        else
            experiment.Summation(settings.summation);
        experiment.FPGAOutputMode(settings.fpga_pixel_output);
    } catch (...) {
        experiment = tmp;
        throw;
    }
}

void LoadDetectorSettings(DiffractionExperiment& experiment, const DetectorSettings &settings) {
    auto tmp = experiment;
    try {
        if (settings.count_time_us > 0)
            experiment.FrameTime(std::chrono::microseconds(settings.frame_time_us),
                                 std::chrono::microseconds(settings.count_time_us));
        else
            experiment.FrameTime(std::chrono::microseconds(settings.frame_time_us));

        experiment.StorageCells(settings.storage_cell_count);
        experiment.UseInternalPacketGenerator(settings.use_internal_packet_generator);

        if (settings.collect_raw_data)
            experiment.Mode(DetectorMode::Raw);
        else
            experiment.Mode(DetectorMode::Conversion);

        experiment.PedestalG0Frames(settings.pedestal_g0_frames);
        experiment.PedestalG1Frames(settings.pedestal_g1_frames);
        experiment.PedestalG2Frames(settings.pedestal_g2_frames);
        if (settings.storage_cell_delay_ns > 0)
            experiment.StorageCellDelay(std::chrono::nanoseconds(settings.storage_cell_delay_ns));
    } catch (...) {
        experiment = tmp;
        throw;
    }
}

JFJochStateMachine::JFJochStateMachine(JFJochServices &in_services, Logger &in_logger)
: services(in_services), logger(in_logger),
data_processing_settings(DiffractionExperiment::DefaultDataProcessingSettings()) {

}

void JFJochStateMachine::ImportPedestalG0(const JFJochReceiverOutput &receiver_output) {
    if (receiver_output.pedestal_result.size() != experiment.GetModulesNum() * experiment.GetStorageCellNumber())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in pedestal output");

    for (int s = 0; s < experiment.GetStorageCellNumber(); s++) {
        for (int module = 0; module < experiment.GetModulesNum(); module++)
            calibration->Pedestal(module, 0, s)
                    = receiver_output.pedestal_result[module + s * experiment.GetModulesNum()];
    }
    SetCalibrationStatistics(calibration->GetModuleStatistics());
}

void JFJochStateMachine::ImportPedestal(const JFJochReceiverOutput &receiver_output, size_t gain_level,
                                        size_t storage_cell) {
    for (int i = 0; i < receiver_output.pedestal_result.size(); i++)
        calibration->Pedestal(i, gain_level, storage_cell) = receiver_output.pedestal_result[i];
    SetCalibrationStatistics(calibration->GetModuleStatistics());
}

void JFJochStateMachine::TakePedestalInternalAll(std::unique_lock<std::mutex> &ul) {
    calibration = std::make_unique<JFCalibration>(experiment);

    if (!gain_calibration.empty()) {
        if (gain_calibration.size() != experiment.GetModulesNum())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in gain files number");
        for (int i = 0; i < gain_calibration.size(); i++)
            calibration->GainCalibration(i) = gain_calibration[i];
    }
    
    cancel_sequence = false;

    logger.Info("Pedestal sequence started");

    try {
        TakePedestalInternalG0(ul);

        for (int i = 0; i < experiment.GetStorageCellNumber(); i++) {
            TakePedestalInternalG1(ul, i);
            TakePedestalInternalG2(ul, i);
        }
    } catch (const std::exception &e) {
        logger.Error("Pedestal sequence error {}", e.what());
        state = JFJochState::Error;
        throw;
    }
    logger.Info("Pedestal sequence done");
}

void JFJochStateMachine::TakePedestalInternalG0(std::unique_lock<std::mutex> &ul) {
    state = JFJochState::Pedestal;
    DiffractionExperiment local_experiment(experiment);
    local_experiment.Mode(DetectorMode::PedestalG0);

    if (local_experiment.GetStorageCellNumber() == 1)
        local_experiment.StorageCellStart(15);
    else
        local_experiment.StorageCellStart(0);

    if (!cancel_sequence && (local_experiment.GetPedestalG0Frames() > 0)) {
        services.Start(local_experiment, *calibration);
        services.Trigger();
        ul.unlock();
        // Allow to cancel/abort during the pedestal data collection
        // Must ensure that while state is Pedestal, nothing can take lock for longer time, to avoid deadlock
        auto pedestal_output = services.Stop(*calibration);
        ul.lock();

        // SetFullMeasurementOutput(pedestal_output);
        ImportPedestalG0(pedestal_output.receiver_output);
    }
    state = JFJochState::Idle;
}

void JFJochStateMachine::TakePedestalInternalG1(std::unique_lock<std::mutex> &ul, int32_t storage_cell) {
    state = JFJochState::Pedestal;
    DiffractionExperiment local_experiment(experiment);
    local_experiment.Mode(DetectorMode::PedestalG1);

    if (local_experiment.GetStorageCellNumber() == 2)
        local_experiment.StorageCellStart((storage_cell + 15) % 16); // one previous
    else
        local_experiment.StorageCellStart(15);

    if (!cancel_sequence && (local_experiment.GetPedestalG1Frames() > 0)) {
        services.Start(local_experiment, *calibration);
        services.Trigger();
        ul.unlock();
        // Allow to cancel/abort during the pedestal data collection
        // Must ensure that while state is Pedestal, nothing can take lock for longer time, to avoid deadlock
        auto pedestal_output = services.Stop(*calibration);
        ul.lock();

        // SetFullMeasurementOutput(pedestal_output);
        ImportPedestal(pedestal_output.receiver_output, 1, storage_cell);
    }
    state = JFJochState::Idle;
}

void JFJochStateMachine::TakePedestalInternalG2(std::unique_lock<std::mutex> &ul, int32_t storage_cell) {
    state = JFJochState::Pedestal;
    DiffractionExperiment local_experiment(experiment);
    local_experiment.Mode(DetectorMode::PedestalG2);

    if (local_experiment.GetStorageCellNumber() == 2)
        local_experiment.StorageCellStart((storage_cell + 15) % 16); // one previous
    else
        local_experiment.StorageCellStart(15);

    if (!cancel_sequence && (local_experiment.GetPedestalG2Frames() > 0)) {
        services.Start(local_experiment, *calibration);
        services.Trigger();
        ul.unlock();
        // Allow to cancel/abort during the pedestal data collection
        // Must ensure that while state is Pedestal, nothing can take lock for longer time, to avoid deadlock
        auto pedestal_output = services.Stop(*calibration);
        ul.lock();

        // SetFullMeasurementOutput(pedestal_output);
        ImportPedestal(pedestal_output.receiver_output, 2, storage_cell);
    }
    state = JFJochState::Idle;
}

void JFJochStateMachine::Initialize() {
    std::unique_lock<std::mutex> ul(m);

    if ((state == JFJochState::Measuring) || (state == JFJochState::Pedestal))
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,  "Cannot initialize during measurement");

    if (detector_setup.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector information not provided");

    logger.Info("Initialize");

    state = JFJochState::Busy;
    ClearMeasurementStatistics();
    try {
        services.On(experiment);
    } catch (...) {
        state = JFJochState::Error;
        throw;
    }
    TakePedestalInternalAll(ul);
}

void JFJochStateMachine::Pedestal() {
    std::unique_lock<std::mutex> ul(m);

    if (state != JFJochState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,"Must be idle to take pedestal");
    TakePedestalInternalAll(ul);
}

void JFJochStateMachine::Trigger() {
    std::unique_lock<std::mutex> ul(m);
    if (state == JFJochState::Measuring)
        services.Trigger();
}

void JFJochStateMachine::Start(const DatasetSettings& settings) {
    std::unique_lock<std::mutex> ul(m);

    if (state != JFJochState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Must be idle to start measurement");

    if (measurement.valid())
        measurement.get(); // In case measurement was running - clear thread

    auto mod_settings = settings;
    SetDatasetDefaults(mod_settings);
    LoadDatasetSettings(experiment, mod_settings);

    ClearAndSetMeasurementStatistics();

    cancel_sequence = false;
    if (experiment.GetStorageCellNumber() == 1)
        experiment.StorageCellStart(15);
    else
        experiment.StorageCellStart(0);

    try {
        state = JFJochState::Busy;
        services.SetDataProcessingSettings(GetDataAnalysisSettings());
        services.Start(experiment, *calibration);

        state = JFJochState::Measuring;
        measurement = std::async(std::launch::async, &JFJochStateMachine::WaitTillMeasurementDone, this);
    } catch (...) {
        state = JFJochState::Error;
        services.Abort();
        throw;
    }
}

void JFJochStateMachine::SetDatasetDefaults(DatasetSettings &settings) {
    if (settings.detector_distance_mm <= 0)
        settings.detector_distance_mm = 100.0;
    if (settings.ntrigger <= 0)
        settings.ntrigger = 1;
}

void JFJochStateMachine::Stop() {
    std::unique_lock<std::mutex> ul(m);

    if (state == JFJochState::Pedestal)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot use the function during pedestal collection");

    c.wait(ul, [&] { return state != JFJochState::Measuring; });

    if (!measurement.valid())
        return; // This is for unlikely condition of two parallel stops
    else
        measurement.get();

    switch (state) {
        case JFJochState::Inactive:
            throw JFJochException(JFJochExceptionCategory::WrongDAQState,"Not initialized");
        case JFJochState::Error:
            throw JFJochException(JFJochExceptionCategory::WrongDAQState,"Detector in error state");
        case JFJochState::Measuring:
        case JFJochState::Busy:
        case JFJochState::Pedestal:
            throw JFJochException(JFJochExceptionCategory::WrongDAQState,"Detector in not expected state to end measurment");
        case JFJochState::Idle:
            break;
    }
}

void JFJochStateMachine::WaitTillMeasurementDone() {
    try {
        auto tmp_output = services.Stop(*calibration);
        SetFullMeasurementOutput(tmp_output);
        {
            std::unique_lock<std::mutex> ul(m);
            state = JFJochState::Idle;
        }
    } catch (...) {
        std::unique_lock<std::mutex> ul(m);
        state = JFJochState::Error;
    }
    c.notify_all();
}

void JFJochStateMachine::Abort() {
    // This is inconsistency in naming - need to solve later
    std::unique_lock<std::mutex> ul(m);
    if ((state == JFJochState::Pedestal) || (state == JFJochState::Measuring)) {
        services.Abort();
        cancel_sequence = true;
    }
}

void JFJochStateMachine::Cancel() {
    // This is inconsistency in naming - need to solve later
    std::unique_lock<std::mutex> ul(m);
    if ((state == JFJochState::Pedestal) || (state == JFJochState::Measuring)) {
        services.Cancel();
        cancel_sequence = true;
    }
}

void JFJochStateMachine::DebugOnly_SetState(JFJochState in_state) {
    std::unique_lock<std::mutex> ul(m);
    state = in_state;
}

void JFJochStateMachine::Deactivate() {
    std::unique_lock<std::mutex> ul(m);
    try {
        if (measurement.valid())
            measurement.get();
        services.Off();
        state = JFJochState::Inactive;
    } catch (...) {
        state = JFJochState::Error;
        throw;
    }
}

JFJochStateMachine::~JFJochStateMachine() {
    try {
        if (measurement.valid())
            measurement.get();
    } catch (...) {}
}

void JFJochStateMachine::SetFullMeasurementOutput(const JFJochServicesOutput &output) {
    std::unique_lock<std::mutex> ul(last_receiver_output_mutex);

    MeasurementStatistics tmp{}; // reset last measurement statistics

    tmp.file_prefix = experiment.GetFilePrefix();
    tmp.detector_width = experiment.GetXPixelsNum();
    tmp.detector_height = experiment.GetYPixelsNum();
    tmp.detector_pixel_depth = experiment.GetPixelDepth();

    tmp.compression_ratio = output.receiver_output.compressed_ratio;
    tmp.collection_efficiency = output.receiver_output.efficiency;
    tmp.images_collected = output.receiver_output.images_sent;
    tmp.cancelled = output.receiver_output.cancelled;
    tmp.max_image_number_sent = output.receiver_output.max_image_number_sent;
    tmp.max_receive_delay = output.receiver_output.max_receive_delay;
    tmp.indexing_rate = output.receiver_output.indexing_rate;
    tmp.bkg_estimate = output.receiver_output.bkg_estimate;

    measurement_statistics = tmp;
}

void JFJochStateMachine::ClearAndSetMeasurementStatistics() {
    std::unique_lock<std::mutex> ul(last_receiver_output_mutex);
    MeasurementStatistics tmp{};

    tmp.file_prefix = experiment.GetFilePrefix();
    tmp.detector_height = experiment.GetXPixelsNum();
    tmp.detector_width = experiment.GetYPixelsNum();
    tmp.detector_pixel_depth = experiment.GetPixelDepth();
    measurement_statistics = tmp;
}

void JFJochStateMachine::ClearMeasurementStatistics() {
    std::unique_lock<std::mutex> ul(last_receiver_output_mutex);
    measurement_statistics.reset();
}

std::optional<MeasurementStatistics> JFJochStateMachine::GetMeasurementStatistics() const {
    std::unique_lock<std::mutex> ul(last_receiver_output_mutex);
    return measurement_statistics;
}

void JFJochStateMachine::LoadMask(const std::vector<uint32_t> &vec, uint32_t bit) {
    std::unique_lock<std::mutex> ul(m);
    if (state == JFJochState::Inactive)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Detector not calibrated");
    if (state != JFJochState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Cannot load mask if detector is not idle");
    calibration->LoadMask(experiment, vec, bit);
}

std::vector<JFCalibrationModuleStatistics> JFJochStateMachine::GetCalibrationStatistics() const {
    std::unique_lock<std::mutex> ul(calibration_statistics_mutex);
    return calibration_statistics;
}

void JFJochStateMachine::SetCalibrationStatistics(const std::vector<JFCalibrationModuleStatistics> &input) {
    std::unique_lock<std::mutex> ul(calibration_statistics_mutex);
    calibration_statistics = input;
}


DetectorSettings JFJochStateMachine::GetDetectorSettings() const {
    std::unique_lock<std::mutex> ul(m);
    DetectorSettings ret;
    ret.frame_time_us = experiment.GetFrameTime().count();
    ret.count_time_us = experiment.GetFrameCountTime().count();
    ret.collect_raw_data = (experiment.GetDetectorMode() != DetectorMode::Conversion);
    ret.use_internal_packet_generator = experiment.IsUsingInternalPacketGen();
    ret.storage_cell_count = experiment.GetStorageCellNumber();
    ret.pedestal_g0_frames = experiment.GetPedestalG0Frames();
    ret.pedestal_g1_frames = experiment.GetPedestalG1Frames();
    ret.pedestal_g2_frames = experiment.GetPedestalG2Frames();
    ret.storage_cell_delay_ns = experiment.GetStorageCellDelay().count();
    return ret;
}

void JFJochStateMachine::SetDetectorSettings(const DetectorSettings &settings) {
    std::unique_lock<std::mutex> ul(m);
    switch (state) {
        case JFJochState::Inactive:
        case JFJochState::Error:
            LoadDetectorSettings(experiment, settings);
            break;
        case JFJochState::Idle:
            LoadDetectorSettings(experiment, settings);
            TakePedestalInternalAll(ul);
            break;
        case JFJochState::Measuring:
        case JFJochState::Busy:
        case JFJochState::Pedestal:
            throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                                  "Cannot change detector set during data collection");
    }
}

DiffractionExperiment &JFJochStateMachine::NotThreadSafe_Experiment() {
    return experiment;
}

BrokerStatus JFJochStateMachine::GetStatus() const {
    BrokerStatus ret{};

    ret.broker_state = state;
    try {
        auto rcv_status = services.GetReceiverStatus();
        ret.progress = rcv_status.progress;
        ret.indexing_rate = rcv_status.indexing_rate;
        ret.receiver_send_buffers_avail = rcv_status.send_buffers_avail;
    } catch (JFJochException &e) {} // ignore exception in getting receiver status (don't really care, e.g. if receiver is down)

    return ret;
}

Plot JFJochStateMachine::GetPlots(const PlotRequest &request) const {
    return services.GetPlots(request);
}

RadialIntegrationProfiles JFJochStateMachine::GetRadialIntegrationProfiles() const {
    return services.GetRadialIntegrationProfiles();
}

void JFJochStateMachine::SetDataProcessingSettings(const DataProcessingSettings &settings) {
    std::unique_lock<std::mutex> ul(data_processing_settings_mutex);
    DiffractionExperiment::CheckDataProcessingSettings(settings);
    data_processing_settings = settings;
    services.SetDataProcessingSettings(settings);
}

DataProcessingSettings JFJochStateMachine::GetDataProcessingSettings() const {
    std::unique_lock<std::mutex> ul(data_processing_settings_mutex);
    return data_processing_settings;
}

DataProcessingSettings JFJochStateMachine::GetDataAnalysisSettings() const {
    std::unique_lock<std::mutex> ul(data_processing_settings_mutex);
    return data_processing_settings;
}

JFJochState JFJochStateMachine::GetState() const {
    return state;
}

void JFJochStateMachine::AddDetectorSetup(const DetectorSetup &setup) {
    std::unique_lock<std::mutex> ul(m);

    if (detector_setup.empty()) {
        experiment.Detector(setup);
        gain_calibration = setup.GetGainCalibration();
        current_detector_setup = 0;
    }
    detector_setup.emplace_back(setup);
}

DetectorList JFJochStateMachine::GetDetectorsList() {
    std::unique_lock<std::mutex> ul(m);
    DetectorList ret;

    for (int i = 0; i < detector_setup.size(); i++) {
        DetectorListElement tmp;
        tmp.description = detector_setup[i].GetDescription();
        tmp.nmodules = detector_setup[i].GetModulesNum();
        tmp.id = i;
        ret.detector.emplace_back(std::move(tmp));
    }

    ret.current_id = current_detector_setup;
    ret.current_description = experiment.GetDetectorDescription();
    return ret;
}

void JFJochStateMachine::SelectDetector(int64_t id) {
    std::unique_lock<std::mutex> ul(m);

    if ((id < 0) || (id >= detector_setup.size()))
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Detector doesn't exist");

    switch (state) {
        case JFJochState::Inactive:
        case JFJochState::Error:
        case JFJochState::Idle:
            try {
                experiment.Detector(detector_setup[id]);
                gain_calibration = detector_setup[id].GetGainCalibration();
                state = JFJochState::Inactive;
                current_detector_setup = id;
            } catch (JFJochException &e) {
                logger.ErrorException(e);
                state = JFJochState::Inactive;
            }
            break;
        case JFJochState::Measuring:
        case JFJochState::Busy:
        case JFJochState::Pedestal:
            throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                                  "Cannot change detector during data collection");
    }
}