#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <limits>
#include <numeric>

namespace fs = std::filesystem;

struct ScanData {
    struct CscanData {
        int Rows;
        int AscanPoints;
        double TsGate;
        double TendGate;
        std::vector<double> X;
        double Cl;
        double Cs;
        double Density;
        double Frequency;
        double Bandwidth;
        char WaveType;
        double BeamAngle;
        double ProbeDiameter;
        double MaxSdtRef;
        double Y;
        int Cols;
        std::vector<std::vector<double>> CscanValues;
        double Wavelength;
        double Nearfield;
        double NearfieldTime;
        double ProbeEllipX;
        double ProbeEllipY;
        double TrueAngle;
        double CalValue;
        std::string DefectType;
        double PosX;
        double PosY;
        double DepthCentre;
        double Diameter;
        double Height;
        double Tilt;
        double MaxSignal;
    };

    CscanData CscanData;
    std::vector<std::vector<double>> AscanValues;
    std::vector<double> timeScale;
};

ScanData datNDTread(const std::map<std::string, std::variant<double, std::string>>& args) {
    ScanData scanData;

    // Parse input arguments
    double TsGate = std::get<double>(args.at("TsGate"));
    double TendGate = std::get<double>(args.at("TendGate"));
    double deltaT = std::get<double>(args.at("deltaT"));
    double minX = std::get<double>(args.at("minX"));
    double maxX = std::get<double>(args.at("maxX"));
    double deltaX = std::get<double>(args.at("deltaX"));
    double backDepth = std::get<double>(args.at("backDepth"));
    double Density = std::get<double>(args.at("Density"));
    double fc = std::get<double>(args.at("fc"));
    double bw = std::get<double>(args.at("bw"));
    char WaveType = std::get<std::string>(args.at("WaveType"))[0];
    double BeamAngle = std::get<double>(args.at("BeamAngle"));
    double d = std::get<double>(args.at("d"));
    std::string baseFolder = std::get<std::string>(args.at("baseFolder"));
    double gain = std::get<double>(args.at("gain"));
    double A0 = std::get<double>(args.at("A0"));
    double alpha = std::get<double>(args.at("alpha"));

    // Get file list in base folder
    std::vector<fs::path> files;
    for (const auto& entry : fs::directory_iterator(baseFolder)) {
        if (entry.path().extension() == ".dat") {
            files.push_back(entry.path());
        }
    }

    // Create AscanValues element
    scanData.CscanData.Rows = files.size();
    scanData.AscanValues.resize(scanData.CscanData.Rows);

    // Import data as columns
    std::vector<double> backWallIdx(files.size());
    for (size_t i = 0; i < files.size(); ++i) {
        std::string filename = files[i].filename().string();
        int col;
        sscanf(filename.c_str(), "pos%d.dat", &col);
        col++;

        std::vector<double> dat;
        std::ifstream file(files[i]);
        double value;
        while (file >> value) {
            dat.push_back(value);
        }

        size_t lenDat = dat.size();
        auto idxm = std::min_element(dat.begin(), dat.begin() + lenDat / 3) - dat.begin();
        auto limIdx = std::min_element(dat.begin() + 2 * lenDat / 3, dat.end()) - dat.begin();

        auto idx2m = std::max_element(dat.begin() + 2 * lenDat / 3, dat.begin() + limIdx) - dat.begin();

        backWallIdx[col - 1] = idx2m - idxm;
        scanData.AscanValues[col - 1].resize(lenDat, 0);
        std::copy(dat.begin() + idxm, dat.end(), scanData.AscanValues[col - 1].begin());
    }

    // Create timeScale element
    scanData.CscanData.AscanPoints = scanData.AscanValues[0].size();
    scanData.timeScale.resize(scanData.CscanData.AscanPoints);
    std::iota(scanData.timeScale.begin(), scanData.timeScale.end(), 0);
    std::transform(scanData.timeScale.begin(), scanData.timeScale.end(), scanData.timeScale.begin(),
                   [deltaT](double x) { return x * deltaT; });

    // Define the time window for simulation
    std::vector<double> t;
    for (double time = TsGate; time <= TendGate; time += deltaT) {
        t.push_back(time);
    }
    size_t nt = t.size();
    size_t idxTsGate = std::lower_bound(t.begin(), t.end(), TsGate) - t.begin();
    size_t idxTendGate = std::lower_bound(t.begin(), t.end(), TendGate) - t.begin();

    // Update scanData with arguments values
    scanData.CscanData.TsGate = TsGate;
    scanData.CscanData.TendGate = TendGate;

    scanData.CscanData.X.resize(scanData.CscanData.Rows);
    std::iota(scanData.CscanData.X.begin(), scanData.CscanData.X.end(), 0);
    std::transform(scanData.CscanData.X.begin(), scanData.CscanData.X.end(), scanData.CscanData.X.begin(),
                   [deltaX](double x) { return x * deltaX; });

    std::vector<double> backWallIdxEnds = {backWallIdx.front(), backWallIdx.back()};
    double Cl = backDepth * 2000 / std::accumulate(backWallIdxEnds.begin(), backWallIdxEnds.end(), 0.0,
                                                   [&](double sum, double idx) { return sum + scanData.timeScale[idx]; });
    scanData.CscanData.Cl = Cl;
    scanData.CscanData.Cs = 0;
    scanData.CscanData.Density = Density;

    scanData.CscanData.Frequency = fc;
    scanData.CscanData.Bandwidth = bw;
    scanData.CscanData.WaveType = WaveType;
    scanData.CscanData.BeamAngle = BeamAngle;
    scanData.CscanData.ProbeDiameter = d;

    // Others elements
    scanData.CscanData.MaxSdtRef = 0;
    scanData.CscanData.Y = 0;
    scanData.CscanData.Cols = 1;
    scanData.CscanData.CscanValues.resize(scanData.CscanData.Rows, std::vector<double>(scanData.CscanData.Cols, 0));

    double c = (WaveType == 'L' || WaveType == 'l') ? scanData.CscanData.Cl : scanData.CscanData.Cs;
    scanData.CscanData.Wavelength = c / (scanData.CscanData.Frequency * 1000);
    scanData.CscanData.Nearfield = std::pow(scanData.CscanData.ProbeDiameter / 2, 2) / scanData.CscanData.Wavelength;
    scanData.CscanData.NearfieldTime = scanData.CscanData.Nearfield / (c / 2000);
    scanData.CscanData.ProbeEllipX = 0;
    scanData.CscanData.ProbeEllipY = 0;
    scanData.CscanData.TrueAngle = scanData.CscanData.BeamAngle;
    scanData.CscanData.CalValue = 0;

    scanData.CscanData.DefectType = "empty";
    scanData.CscanData.PosX = 0;
    scanData.CscanData.PosY = 0;
    scanData.CscanData.DepthCentre = 0;
    scanData.CscanData.Diameter = 0;
    scanData.CscanData.Height = 0;
    scanData.CscanData.Tilt = 0;

    // Adjust A-scan amplitude
    if (std::isinf(gain)) {
        double MaxSignal = 0;
        for (const auto& ascan : scanData.AscanValues) {
            MaxSignal = std::max(MaxSignal, *std::max_element(ascan.begin(), ascan.end(), [](double a, double b) { return std::abs(a) < std::abs(b); }));
        }
        for (auto& ascan : scanData.AscanValues) {
            std::transform(ascan.begin(), ascan.end(), ascan.begin(), [MaxSignal](double x) { return x / MaxSignal; });
        }
    } else {
        double gainFactor = std::pow(10, gain / 20);
        for (auto& ascan : scanData.AscanValues) {
            std::transform(ascan.begin(), ascan.end(), ascan.begin(), [gainFactor](double x) { return (0.5 / 2048 / gainFactor) * x; });
        }
    }

    // Apply TGV
    std::vector<double> zz(scanData.timeScale.size());
    std::transform(scanData.timeScale.begin(), scanData.timeScale.end(), zz.begin(), [Cl](double t) { return t * Cl / 2000; });
    std::vector<double> ffc(zz.size());
    std::transform(zz.begin(), zz.end(), ffc.begin(), [A0, alpha](double z) { return std::max(1.0, (1 / A0) * std::exp(alpha * z)); });

    for (auto& ascan : scanData.AscanValues) {
        std::transform(ascan.begin(), ascan.end(), ffc.begin(), ascan.begin(), std::multiplies<double>());
    }

    // Calculate elements after the simulation session
    scanData.CscanData.MaxSignal = 0;
    for (const auto& ascan : scanData.AscanValues) {
        scanData.CscanData.MaxSignal = std::max(scanData.CscanData.MaxSignal, *std::max_element(ascan.begin(), ascan.end(), [](double a, double b) { return std::abs(a) < std::abs(b); }));
    }

    for (size_t i = 0; i < scanData.CscanData.Rows; ++i) {
        double maxAbs = *std::max_element(scanData.AscanValues[i].begin(), scanData.AscanValues[i].end(), [](double a, double b) { return std::abs(a) < std::abs(b); });
        scanData.CscanData.CscanValues[i][0] = 10 * std::log10(maxAbs);
    }

    return scanData;
}