/**
 * negfix_cpp - High performance C++ port of negfix8
 * Expects Linear Rec 709 16-bit TIFF Input.
 *
 * Usage:
 * ./negfix -i input.tif -o output.tif [flags]
 */

#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <opencv2/opencv.hpp>
#include <omp.h>

// Default settings from original script
struct Settings {
    double gamma = 2.15;
    bool contrastStretch = false;
    bool mirror = false;
    int binning = 1;
    int saturationBoost = 0;  // 0 = disabled, 1-100 = boost percentage
    double exposureStops = 0.0;  // 0 = disabled, range -10 to +10
    std::string profileLoadPath = "";
    std::string profileSavePath = "";
    std::string inputPath = "";
    std::string outputPath = "";
};

struct Profile {
    double minR = 0, minG = 0, minB = 0;
    double maxR = 0, maxG = 0, maxB = 0;
    double gammaG = 0, gammaB = 0;
    double offset = 0; // Calculated offset for black point
    bool valid = false;
};

// Helper to parse arguments
bool parseArgs(int argc, char** argv, Settings& settings) {
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "-i" && i + 1 < argc) settings.inputPath = argv[++i];
        else if (arg == "-o" && i + 1 < argc) settings.outputPath = argv[++i];
        else if (arg == "-g" && i + 1 < argc) settings.gamma = std::stod(argv[++i]);
        else if (arg == "-cs") settings.contrastStretch = true;
        else if (arg == "-m") settings.mirror = true;
        else if (arg == "-sat") {
            // Check if next argument is a number (optional value)
            if (i + 1 < argc) {
                try {
                    int val = std::stoi(argv[i + 1]);
                    if (val >= 0 && val <= 100) {
                        settings.saturationBoost = val;
                        ++i;  // consume the value argument
                    } else {
                        settings.saturationBoost = 20;  // default if out of range
                    }
                } catch (...) {
                    // Next arg is not a number, use default
                    settings.saturationBoost = 20;
                }
            } else {
                settings.saturationBoost = 20;  // default
            }
        }
        else if (arg == "-exp" && i + 1 < argc) {
            double val = std::stod(argv[++i]);
            if (val < -10.0 || val > 10.0) {
                std::cerr << "Error: -exp value must be between -10 and +10 stops\n";
                exit(1);
            }
            settings.exposureStops = val;
        }
        else if (arg == "-r" && i + 1 < argc) settings.binning = std::stoi(argv[++i]);
        else if (arg == "-c" && i + 1 < argc) settings.profileSavePath = argv[++i];
        else if (arg == "-u" && i + 1 < argc) settings.profileLoadPath = argv[++i];
        else if (arg == "-h" || arg == "--help") return false;
    }
    return !settings.inputPath.empty();
}

void printUsage() {
    std::cout << "Usage: negfix_cpp -i input.tif [-o output.tif] [flags]\n"
              << "Flags:\n"
              << "  -g <val>    Set gamma (default 2.15)\n"
              << "  -cs         Enable contrast stretching\n"
              << "  -exp <N>    Exposure adjustment in stops (-10 to +10)\n"
              << "  -sat [N]    Saturation boost, N=0-100 (default 20)\n"
              << "  -m          Mirror image horizontally\n"
              << "  -r <N>      Binning (scale down by 1/N)\n"
              << "  -c <file>   Create profile and save to file (exits after)\n"
              << "  -u <file>   Use existing profile\n";
}

// Analyze image to find Min/Max (Film Base and Highlights)
// Simulates: -shave 10x10 -blur 3x3 then finding min/max
Profile analyzeImage(const cv::Mat& img) {
    Profile prof;
    
    // Create a copy for analysis to avoid modifying original yet
    cv::Mat work;
    img.copyTo(work);

    // 1. Shave 10 pixels
    int shave = 10;
    if (work.cols > 20 && work.rows > 20) {
        cv::Rect roi(shave, shave, work.cols - (shave*2), work.rows - (shave*2));
        work = work(roi);
    }

    // 2. Blur 3x3
    cv::blur(work, work, cv::Size(3, 3));

    // 3. Split channels
    std::vector<cv::Mat> channels(3);
    cv::split(work, channels);

    // 4. Find Min/Max
    // Note: OpenCV loads 16bit TIFF as CV_16U.
    // In original script:
    // MINIMA (Darkest in file) corresponds to Scene HIGHLIGHTS (Densest negative)
    // MAXIMA (Brightest in file) corresponds to Scene SHADOWS (Clear film base)
    
    double minVal, maxVal;
    
    // Blue (channel 0 in OpenCV)
    cv::minMaxLoc(channels[0], &minVal, &maxVal);
    prof.minB = minVal; prof.maxB = maxVal;

    // Green (channel 1)
    cv::minMaxLoc(channels[1], &minVal, &maxVal);
    prof.minG = minVal; prof.maxG = maxVal;

    // Red (channel 2)
    cv::minMaxLoc(channels[2], &minVal, &maxVal);
    prof.minR = minVal; prof.maxR = maxVal;

    // Validate
    if (prof.minR == 0 || prof.maxR == 0) {
        std::cerr << "Error: Image stats invalid (0 values found). scan might be empty.\n";
        exit(1);
    }

    prof.valid = true;
    return prof;
}

void saveProfile(const std::string& path, const Profile& p, double userGamma) {
    std::ofstream out(path);
    // Format mimics original script: minR minG minB offset gammaG gammaB gammaGlobal
    // Calculate the derived values to store
    
    // Replicating original math:
    // Offset calculation from script:
    // ((minR / maxR)^(1/gamma)) * QuantumRange * 0.95
    double quantum = 65535.0;
    double offset = std::pow(p.minR / p.maxR, 1.0/userGamma) * quantum * 0.95;

    // Gammas
    double gG = std::log(p.maxG / p.minG) / std::log(p.maxR / p.minR);
    double gB = std::log(p.maxB / p.minB) / std::log(p.maxR / p.minR);

    out << p.minR << " " << p.minG << " " << p.minB << " " 
        << offset << " " << gG << " " << gB << " " << userGamma << "\n";
    out.close();
    std::cout << "Profile saved to " << path << "\n";
}

Profile loadProfile(const std::string& path) {
    Profile p;
    std::ifstream in(path);
    if (!in.is_open()) {
        std::cerr << "Could not open profile: " << path << "\n";
        exit(1);
    }
    double dummyGamma;
    in >> p.minR >> p.minG >> p.minB >> p.offset >> p.gammaG >> p.gammaB >> dummyGamma;
    p.valid = true;
    return p;
}

// Generate LUTs for blazing fast application
void generateLUTs(std::vector<uint16_t>& lutR, std::vector<uint16_t>& lutG, std::vector<uint16_t>& lutB, 
                  const Profile& p, double gamma) {
    
    lutR.resize(65536);
    lutG.resize(65536);
    lutB.resize(65536);

    // Calculate Gamma corrections if we analyzed image fresh
    // If loaded from profile, these are already in p.gammaG/B
    double gG, gB, offset;

    if (p.offset == 0 && p.gammaG == 0) {
        // Just analyzed, calculate derived values
        gG = std::log(p.maxG / p.minG) / std::log(p.maxR / p.minR);
        gB = std::log(p.maxB / p.minB) / std::log(p.maxR / p.minR);
        offset = std::pow(p.minR / p.maxR, 1.0/gamma) * 65535.0 * 0.95;
    } else {
        // Loaded from file
        gG = p.gammaG;
        gB = p.gammaB;
        offset = p.offset;
    }

    double invGamma = 1.0 / gamma;

    // OpenMP loop for LUT generation
    #pragma omp parallel for
    for (int i = 0; i < 65536; ++i) {
        double u = (double)i;
        if (u < 1.0) u = 1.0; // avoid div by zero

        // 1. Inversion Logic: Base / Input
        // Original Script: R_out = MinR / Input
        double r_norm = p.minR / u;
        double g_norm = p.minG / u;
        double b_norm = p.minB / u;

        // 2. Channel Gamma Alignment
        // R is reference (gamma 1 relative to itself)
        // G and B: ImageMagick's -gamma X applies pow(value, 1/X)
        // So we need pow(g_norm, 1/gG) not pow(g_norm, gG)
        double r_aligned = r_norm;
        double g_aligned = std::pow(g_norm, 1.0/gG);
        double b_aligned = std::pow(b_norm, 1.0/gB);

        // 3. Global Gamma Application
        double r_final = std::pow(r_aligned, invGamma);
        double g_final = std::pow(g_aligned, invGamma);
        double b_final = std::pow(b_aligned, invGamma);

        // 4. Scale to QuantumRange (assume internal math was normalized to 0..1? 
        // No, min/u results in small numbers if u is large.
        // Wait, if u = minR, r_norm = 1.0. If u = maxR, r_norm = small.
        // The power functions preserve the 0..1 range mostly.
        // We need to scale up to 65535 BEFORE subtraction?
        // Let's trace unit:
        // MinR is ~300. u is ~300. Ratio = 1.
        // MinR is ~300. u is ~60000. Ratio = 0.005.
        // So we are in 0..1 range effectively.
        
        r_final *= 65535.0;
        g_final *= 65535.0;
        b_final *= 65535.0;

        // 5. Subtract Offset
        r_final -= offset;
        g_final -= offset;
        b_final -= offset;

        // Clamp
        lutR[i] = cv::saturate_cast<uint16_t>(r_final);
        lutG[i] = cv::saturate_cast<uint16_t>(g_final);
        lutB[i] = cv::saturate_cast<uint16_t>(b_final);
    }
}

int main(int argc, char** argv) {
    Settings settings;
    if (!parseArgs(argc, argv, settings)) {
        printUsage();
        return 0;
    }

    // Load Image
    std::cout << "Loading " << settings.inputPath << "..." << std::endl;
    cv::Mat img = cv::imread(settings.inputPath, cv::IMREAD_UNCHANGED);
    if (img.empty()) {
        std::cerr << "Failed to load image.\n";
        return 1;
    }

    if (img.depth() != CV_16U) {
        std::cerr << "Error: Input must be 16-bit.\n";
        return 1;
    }

    // Determine Profile
    Profile prof;
    if (!settings.profileLoadPath.empty()) {
        prof = loadProfile(settings.profileLoadPath);
        std::cout << "Loaded profile " << settings.profileLoadPath << "\n";
    } else {
        std::cout << "Analyzing image profile..." << std::endl;
        prof = analyzeImage(img);
        
        if (!settings.profileSavePath.empty()) {
            saveProfile(settings.profileSavePath, prof, settings.gamma);
            return 0; // Exit after creating profile as per script logic behavior
        }
    }

    // Binning (Resize)
    if (settings.binning > 1) {
        cv::resize(img, img, cv::Size(img.cols/settings.binning, img.rows/settings.binning), 0, 0, cv::INTER_AREA);
    }

    // Generate LUTs
    std::cout << "Generating Lookup Tables..." << std::endl;
    std::vector<uint16_t> lutR, lutG, lutB;
    generateLUTs(lutR, lutG, lutB, prof, settings.gamma);

    // Apply LUTs
    std::cout << "Processing pixels..." << std::endl;
    
    // If mono, process differently, but assuming Color input for Negfix logic
    int rows = img.rows;
    int cols = img.cols;

    if (img.channels() == 3) {
        // Parallel Loop over rows
        #pragma omp parallel for
        for (int r = 0; r < rows; ++r) {
            uint16_t* ptr = img.ptr<uint16_t>(r);
            for (int c = 0; c < cols; ++c) {
                // OpenCV is BGR by default
                uint16_t b = ptr[3*c + 0];
                uint16_t g = ptr[3*c + 1];
                uint16_t r_val = ptr[3*c + 2];

                ptr[3*c + 0] = lutB[b];
                ptr[3*c + 1] = lutG[g];
                ptr[3*c + 2] = lutR[r_val];
            }
        }
    } else {
        // Monochrome case
        #pragma omp parallel for
        for (int r = 0; r < rows; ++r) {
            uint16_t* ptr = img.ptr<uint16_t>(r);
            for (int c = 0; c < cols; ++c) {
                 ptr[c] = lutG[ptr[c]]; // Use Green channel logic for B&W usually
            }
        }
    }

    // Optional: Exposure adjustment (applied before saturation)
    if (settings.exposureStops != 0.0) {
        std::cout << "Applying Exposure Adjustment (" << std::showpos << settings.exposureStops 
                  << std::noshowpos << " stops)..." << std::endl;
        
        // multiplier = 2^stops: +1 stop = 2x, -1 stop = 0.5x
        double multiplier = std::pow(2.0, settings.exposureStops);
        
        #pragma omp parallel for
        for (int r = 0; r < rows; ++r) {
            uint16_t* ptr = img.ptr<uint16_t>(r);
            for (int c = 0; c < cols * img.channels(); ++c) {
                ptr[c] = cv::saturate_cast<uint16_t>(ptr[c] * multiplier);
            }
        }
    }

    // Optional: Saturation boost using Rec.709 luminance method
    if (settings.saturationBoost > 0) {
        std::cout << "Applying Saturation Boost (" << settings.saturationBoost << "%)..." << std::endl;
        
        // sat_factor: 1.0 = no change, 1.2 = 20% boost, 2.0 = 100% boost
        double sat_factor = 1.0 + (settings.saturationBoost / 100.0);
        
        // Rec.709 luma coefficients (matches input colorspace)
        const double luma_r = 0.2126;
        const double luma_g = 0.7152;
        const double luma_b = 0.0722;
        
        #pragma omp parallel for
        for (int r = 0; r < rows; ++r) {
            uint16_t* ptr = img.ptr<uint16_t>(r);
            for (int c = 0; c < cols; ++c) {
                // OpenCV is BGR
                double b = ptr[3*c + 0];
                double g = ptr[3*c + 1];
                double r_val = ptr[3*c + 2];
                
                // Calculate Rec.709 luminance
                double luma = luma_r * r_val + luma_g * g + luma_b * b;
                
                // Apply saturation: new = luma + sat_factor * (original - luma)
                double new_r = luma + sat_factor * (r_val - luma);
                double new_g = luma + sat_factor * (g - luma);
                double new_b = luma + sat_factor * (b - luma);
                
                ptr[3*c + 0] = cv::saturate_cast<uint16_t>(new_b);
                ptr[3*c + 1] = cv::saturate_cast<uint16_t>(new_g);
                ptr[3*c + 2] = cv::saturate_cast<uint16_t>(new_r);
            }
        }
    }

    // Optional: Contrast Stretch
    if (settings.contrastStretch) {
        std::cout << "Contrast Stretching..." << std::endl;
        // Simple MinMax normalization
        cv::normalize(img, img, 0, 65535, cv::NORM_MINMAX);
    }

    // Optional: Mirror
    if (settings.mirror) {
        cv::flip(img, img, 1);
    }

    // Save
    std::string out = settings.outputPath.empty() ? "output.tif" : settings.outputPath;
    std::cout << "Saving to " << out << "..." << std::endl;
    
    // Set compression params for TIFF
    std::vector<int> params;
    params.push_back(cv::IMWRITE_TIFF_COMPRESSION);
    params.push_back(5); // LZW (equivalent to Zip roughly in OpenCV terms, usually 5 or 1)

    cv::imwrite(out, img, params);
    std::cout << "Done." << std::endl;

    return 0;
}