#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <optional>
#include <filesystem> // For directory creation and path manipulation
#include <torch/torch.h>
#include <torch/script.h> // For torch::save and torch::load
#include <ATen/Context.h> // Required for globalContext
// #include <opencv2/opencv.hpp> // REMOVED
#include <algorithm> // For std::find

// Project headers
#include "../cimp/resnet/resnet.h"
#include "../cimp/classifier/classifier.h"
#include "../cimp/bb_regressor/bb_regressor.h"

namespace fs = std::filesystem;

// Helper to ensure a directory exists
void ensure_dir_exists(const std::string& path_str) {
    fs::path dir_path(path_str);
    if (!fs::exists(dir_path)) {
        if (fs::create_directories(dir_path)) {
            std::cout << "Created directory: " << dir_path.string() << std::endl;
        } else {
            std::cerr << "Error: Could not create directory: " << dir_path.string() << std::endl;
            // Optionally throw an exception or exit
        }
    }
}

// Helper to load a tensor
torch::Tensor load_tensor_from_file(const std::string& file_path, torch::Device device) {
    std::cout << "  Loading tensor from: " << file_path << std::endl;
    torch::Tensor tensor;
    torch::load(tensor, file_path);
    return tensor.to(device);
}

// Helper to save a tensor
void save_tensor_to_file(const torch::Tensor& tensor, const std::string& file_path) {
    try {
        torch::save(tensor.cpu(), file_path);
    } catch (const c10::Error& e) {
        std::cerr << "Error saving tensor to " << file_path << ": " << e.what() << std::endl;
    }
}

// Helper to generate or load a dummy tensor
torch::Tensor get_or_generate_tensor(
    const std::string& file_path, 
    const std::vector<int64_t>& shape, 
    torch::Device device, 
    bool force_generate = false) {
    if (!force_generate && fs::exists(file_path)) {
        return load_tensor_from_file(file_path, device);
    } else {
        std::cout << "  Generating dummy tensor for: " << file_path << " with shape [";
        for(size_t i=0; i < shape.size(); ++i) std::cout << shape[i] << (i == shape.size()-1 ? "" : ", ");
        std::cout << "]" << std::endl;
        auto tensor = torch::randn(shape, torch::TensorOptions().device(device).dtype(torch::kFloat32));
        save_tensor_to_file(tensor, file_path); // Save the generated one for future use
        return tensor;
    }
}

int main(int argc, char* argv[]) {
    if (argc < 4) {
        // Original usage message for consistency with how run_tests.sh calls it
        std::cerr << "Usage: " << argv[0] << " <base_input_dir_path> <base_output_dir_path> <num_samples> [model_to_test]" << std::endl;
        return 1;
    }

    // --- Set Global PyTorch/cuDNN Flags ---
    at::globalContext().setDeterministicCuDNN(true);
    at::globalContext().setBenchmarkCuDNN(false); // Explicitly set benchmark to false
    std::cout << "C++: Set globalContext DeterministicCuDNN=true, BenchmarkCuDNN=false" << std::endl;
    // ---

    // Original argument parsing structure
    std::string base_input_dir = argv[1];
    std::string base_output_dir = argv[2];
    int num_samples = 0;
    try {
        num_samples = std::stoi(argv[3]);
    } catch (const std::exception& e) {
        std::cerr << "Error: Invalid number of samples in argv[3]='" << argv[3] << "': " << e.what() << std::endl;
        return 1;
    }
    
    std::string model_to_test = "all"; // Default to all
    if (argc > 4) {
        model_to_test = argv[4];
    }

    if (num_samples <= 0) {
        std::cout << "Number of samples is " << num_samples << ". Exiting." << std::endl;
        return 0;
    }
    std::cout << "Running for " << num_samples << " samples." << std::endl;
    std::cout << "Input directory: " << base_input_dir << std::endl;
    std::cout << "Output directory: " << base_output_dir << std::endl;

    torch::DeviceType device_type;
    if (torch::cuda::is_available()) {
        std::cout << "CUDA available! Running on GPU." << std::endl;
        device_type = torch::kCUDA; // REVERTED to kCUDA
    } else {
        std::cout << "CUDA not available. Running on CPU." << std::endl;
        device_type = torch::kCPU;
    }
    torch::Device device(device_type);

    // --- Initialize Models (once) ---
    std::cout << "--- Initializing Models ---" << std::endl;
    std::string resnet_weights_path = "../exported_weights/backbone_regenerated";
    std::string bb_reg_weights_path = "../exported_weights/bb_regressor";
    std::string classifier_weights_path = "../exported_weights/classifier";
    std::optional<cimp::resnet::ResNet> resnet_model_opt;
    // Use TORCH_MODULE wrapper types directly
    std::optional<Classifier> classifier_model_opt_wrapped;
    std::optional<BBRegressor> bb_regressor_model_opt_wrapped;
    bool models_initialized_ok = true;

    try {
        std::vector<std::string> output_layers_resnet = {
            "conv1_output", "bn1_output", "relu1_output", "maxpool_output",
            "layer1", "layer2", "layer3", "layer4", "features", 
            "layer1_0_shortcut_output", "layer1_0_block_output",
            "debug_resnet_conv1_output_for_bn1_input",
            // New BN1 intermediate outputs
            "bn1_centered_x", "bn1_variance_plus_eps", "bn1_inv_std", "bn1_normalized_x"
        };
        resnet_model_opt.emplace(cimp::resnet::resnet50(resnet_weights_path, output_layers_resnet, device)); 
        (*resnet_model_opt)->to(device);
        (*resnet_model_opt)->eval();
        std::cout << "ResNet-50 initialized successfully." << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "CRITICAL: Error initializing ResNet-50: " << e.what() << std::endl;
        models_initialized_ok = false; 
    }
    
    if (models_initialized_ok) { 
        try {
            std::cout << "Attempting to initialize BBRegressor with weights from: " << bb_reg_weights_path << std::endl;
            // BBRegressor class is in global namespace
            bb_regressor_model_opt_wrapped.emplace(BBRegressor(bb_reg_weights_path, device));
            (*bb_regressor_model_opt_wrapped).to(device);
            (*bb_regressor_model_opt_wrapped).eval();
            std::cout << "BBRegressor initialized successfully." << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "CRITICAL: Error initializing BBRegressor: " << e.what() << std::endl;
            models_initialized_ok = false;
        }
    }

    if (models_initialized_ok) { 
        try {
            std::cout << "Attempting to initialize Classifier with weights from: " << classifier_weights_path << std::endl;
            // Classifier class is in global namespace
            classifier_model_opt_wrapped.emplace(Classifier(classifier_weights_path, device));
            (*classifier_model_opt_wrapped).to(device);
            // (*classifier_model_opt_wrapped).eval(); // Classifier class does not have an eval() method
            std::cout << "Classifier initialized successfully." << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "CRITICAL: Error initializing Classifier: " << e.what() << std::endl;
            models_initialized_ok = false;
        }
    }

    if (!models_initialized_ok) {
        std::cerr << "Model initialization failed. Aborting tests." << std::endl;
        return 1;
    }
    std::cout << "All models initialized." << std::endl;

    // --- Define and Announce Output Directories Before Loop ---
    fs::path resnet_out_dir_main = fs::path(base_output_dir) / "resnet";
    fs::path clf_out_dir_main = fs::path(base_output_dir) / "classifier";
    fs::path bb_reg_out_dir_main = fs::path(base_output_dir) / "bb_regressor";

    std::cout << "C++ Model Output Directories Check:" << std::endl; // Changed text slightly for new check
    std::cout << "  Expected ResNet output: " << resnet_out_dir_main.string() << std::endl;
    std::cout << "  Expected Classifier output: " << clf_out_dir_main.string() << std::endl;
    std::cout << "  Expected BB Regressor output: " << bb_reg_out_dir_main.string() << std::endl;

    // --- Main Loop for Samples ---
    for (int i = 0; i < num_samples; ++i) {
        std::cout << "--- Processing Sample " << i << " ---" << std::endl;
        std::string sample_suffix = "sample_" + std::to_string(i);

        // Define input file paths
        fs::path common_input_path = fs::path(base_input_dir) / "common";
        ensure_dir_exists(common_input_path.string()); 

        std::string image_file = (common_input_path / (sample_suffix + "_image.pt")).string();
        std::string bb_file = (common_input_path / (sample_suffix + "_bb.pt")).string();
        std::string proposals_file = (common_input_path / (sample_suffix + "_proposals.pt")).string();

        // Define output directories
        fs::path resnet_out_dir = fs::path(base_output_dir) / "resnet";
        fs::path clf_out_dir = fs::path(base_output_dir) / "classifier";
        fs::path bb_reg_out_dir = fs::path(base_output_dir) / "bb_regressor";
        ensure_dir_exists(resnet_out_dir.string());
        ensure_dir_exists(clf_out_dir.string());
        ensure_dir_exists(bb_reg_out_dir.string());

        bool force_generate_dummy = (i == 0); 

        torch::Tensor image_tensor = get_or_generate_tensor(image_file, {1, 3, 256, 256}, device, force_generate_dummy && !fs::exists(image_file));
        torch::Tensor bb_tensor = get_or_generate_tensor(bb_file, {1, 4}, device, force_generate_dummy && !fs::exists(bb_file)); 
        torch::Tensor proposals_tensor = get_or_generate_tensor(proposals_file, {1, 10, 4}, device, force_generate_dummy && !fs::exists(proposals_file)); 

        // --- Normalize the image tensor --- START
        if (image_tensor.dim() == 3) { // Add batch dim if not present
            image_tensor = image_tensor.unsqueeze(0);
        }
        // Ensure it's float (already done by get_or_generate_tensor if is_image=true)
        image_tensor = image_tensor / 255.0f; // Add division by 255.0

        torch::Tensor mean = torch::tensor({0.485, 0.456, 0.406}, image_tensor.options()).view({1, -1, 1, 1});
        torch::Tensor std_dev = torch::tensor({0.229, 0.224, 0.225}, image_tensor.options()).view({1, -1, 1, 1});
        image_tensor.sub_(mean).div_(std_dev);
        // --- Normalize the image tensor --- END

        // Save C++ preprocessed input tensor
        fs::path cpp_preprocessed_save_path = resnet_out_dir / ("sample_" + std::to_string(i) + "_image_preprocessed_cpp.pt");
        save_tensor_to_file(image_tensor, cpp_preprocessed_save_path.string());
        std::cout << "Saved C++ preprocessed image for sample " << i << " to " << cpp_preprocessed_save_path.string() << std::endl;

        // 1. ResNet Processing
        std::map<std::string, torch::Tensor> resnet_outputs;
        if (resnet_model_opt.has_value()){ 
            try {
                std::cout << "Processing ResNet for sample " << i << std::endl;
                resnet_outputs = (*resnet_model_opt)->forward(image_tensor);
                // Save all defined output layers for ResNet for comparison
                // The output_layers used during construction was: {"layer1", "layer2", "layer3", "layer4", "features", "layer1_0_shortcut_output"}
                for (const auto& pair : resnet_outputs) {
                    std::string output_name = pair.first;
                    std::string filename_key = output_name;
                    std::replace(filename_key.begin(), filename_key.end(), '.', '_'); // Ensure dots are replaced for all keys

                    fs::path output_path = resnet_out_dir / ("sample_" + std::to_string(i) + "_" + filename_key + ".pt");
                    save_tensor_to_file(pair.second, output_path.string());
                    std::cout << "  Saving tensor to: " << output_path.string() << " shape: " << pair.second.sizes() << std::endl;
                }
                std::cout << "ResNet processing done for sample " << i << std::endl;
            } catch (const std::exception& e) {
                std::cerr << "Error during ResNet processing for sample " << i << ": " << e.what() << std::endl;
                // continue; // Decide if to skip to next sample or let it fail further down
            }
        } else {
            std::cerr << "Skipping ResNet processing for sample " << i << " as model was not initialized." << std::endl;
            // If ResNet didn't initialize, resnet_outputs will be empty. Downstream models need to handle this.
        }
        
        // 2. Classifier Processing
        if (classifier_model_opt_wrapped.has_value() && resnet_outputs.count("layer3")) { 
            try {
                std::cout << "Processing Classifier for sample " << i << std::endl;
                auto clf_input = resnet_outputs["layer3"].clone(); 
                torch::Tensor clf_features = (*classifier_model_opt_wrapped).extract_features(clf_input);
                save_tensor_to_file(clf_features, (clf_out_dir / (sample_suffix + "_features.pt")).string());
                std::cout << "Classifier processing done for sample " << i << std::endl;
            } catch (const std::exception& e) {
                std::cerr << "Error during Classifier processing for sample " << i << ": " << e.what() << std::endl;
            }
        } else {
             std::cerr << "Skipping Classifier for sample " << i << " (model not init or ResNet layer3 input missing)." << std::endl;
        }
        
        // 3. BBRegressor Processing
        if (bb_regressor_model_opt_wrapped.has_value() && resnet_outputs.count("layer2") && resnet_outputs.count("layer3")) { 
            try {
                std::cout << "Processing BBRegressor for sample " << i << std::endl;
                // Remove debug_get_conv3_1t_output and debug_get_conv4_1t_output calls (they do not exist)
                // The correct BBRegressor logic is:
                std::vector<torch::Tensor> backbone_feats_for_bb = {
                    resnet_outputs["layer2"].clone(), 
                    resnet_outputs["layer3"].clone()  
                };
                
                std::vector<torch::Tensor> iou_feats = (*bb_regressor_model_opt_wrapped).get_iou_feat(backbone_feats_for_bb);
                if (iou_feats.size() >= 1) save_tensor_to_file(iou_feats[0], (bb_reg_out_dir / (sample_suffix + "_iou_feat0.pt")).string());
                if (iou_feats.size() >= 2) save_tensor_to_file(iou_feats[1], (bb_reg_out_dir / (sample_suffix + "_iou_feat1.pt")).string());

                std::vector<torch::Tensor> mod_vectors = (*bb_regressor_model_opt_wrapped).get_modulation(backbone_feats_for_bb, bb_tensor);
                if (mod_vectors.size() >= 1) save_tensor_to_file(mod_vectors[0], (bb_reg_out_dir / (sample_suffix + "_mod_vec0.pt")).string());
                if (mod_vectors.size() >= 2) save_tensor_to_file(mod_vectors[1], (bb_reg_out_dir / (sample_suffix + "_mod_vec1.pt")).string());

                if (!iou_feats.empty() && !mod_vectors.empty()) {
                    torch::Tensor iou_scores = (*bb_regressor_model_opt_wrapped).predict_iou(mod_vectors, iou_feats, proposals_tensor);
                    save_tensor_to_file(iou_scores, (bb_reg_out_dir / (sample_suffix + "_iou_scores.pt")).string());
                } else {
                    std::cerr << "  Skipping BBRegressor predict_iou for sample " << i << " (iou_feats or mod_vectors empty)." << std::endl;
                }
                std::cout << "BBRegressor processing done for sample " << i << std::endl;
            } catch (const std::exception& e) {
                std::cerr << "Error during BBRegressor processing for sample " << i << ": " << e.what() << std::endl;
            }
        } else {
            std::cerr << "Skipping BBRegressor for sample " << i << " (model not init or ResNet inputs missing)." << std::endl;
        }
    }

    std::cout << "File-based test models run completed for " << num_samples << " samples." << std::endl;
    return 0;
}