#include "map_loader.hpp"
#include <sstream>
#include <stdexcept>
#include <utility>

namespace maploader {

// --- model forward pass ----------------------------------------------------------------

/*
Summary:
  Construct a per-axis inference model with shared trunk and two output heads.
Args:
  in_features: Number of input features per sample.
  hidden: Hidden layer widths.
Returns:
  None.
Side Effects:
  Registers Torch submodules (`body`, `order`, `modes`).
Raises:
  None.
Preconditions:
  `in_features > 0` and hidden widths are positive.
*/
ShaperNetImpl::ShaperNetImpl(int in_features, std::vector<int> hidden) {
    body = register_module("body", torch::nn::Sequential());
    int last = in_features;
    for (int h : hidden) {
        body->push_back(torch::nn::Linear(last, h));
        body->push_back(torch::nn::ReLU());
        last = h;
    }
    order_head = register_module("order", torch::nn::Linear(last, 1));  // binary
    mode_head  = register_module("modes", torch::nn::Linear(last, 4));  // [wn1, lnζ1, wn2, lnζ2]
}

/*
Summary:
  Run one forward pass and return order probability plus mode parameters.
Args:
  x: Input tensor shaped `[N, in_features]`.
Returns:
  `std::pair<torch::Tensor, torch::Tensor>` containing `(order_prob [N,1], mode_params [N,4])`.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  `x` has feature width equal to the model input width.
*/
std::pair<torch::Tensor, torch::Tensor>
ShaperNetImpl::forward(torch::Tensor x) {
    x = body->forward(x);
    auto order_logits = order_head->forward(x);
    auto order = torch::sigmoid(order_logits);
    auto half = torch::full_like(order, 0.5f);
    order = torch::where(order.isnan(), half, order);
    order = order.clamp(1e-6f, 1.0f - 1e-6f);
    auto modes = mode_head->forward(x);
    return {order, modes};
}

// --- driver ---------------------------------------------------------------

MapFitter::MapFitter(int axes, int input_features, std::vector<int> hidden)
    : axes_(axes), in_features_(input_features), hidden_(std::move(hidden)) {
    models_.reserve(axes_);
    for (int i = 0; i < axes_; ++i) {
        models_.emplace_back(in_features_, hidden_);
    }
}

/*
Summary:
  Infer mode/order outputs for one axis.
Args:
  axis: Axis index in `[0, axes_)`.
  feature: Feature tensor for that axis.
Returns:
  `std::pair<torch::Tensor, torch::Tensor>` axis outputs.
Side Effects:
  Switches selected model to evaluation mode.
Raises:
  std::out_of_range: If `axis` is invalid.
Preconditions:
  `feature` shape matches model expectations.
*/
std::pair<torch::Tensor, torch::Tensor>
MapFitter::infer(int axis, const torch::Tensor &feature) {
    auto X = feature.contiguous().to(torch::kFloat32);
    if (scripted_) {
        auto &module = scripted_models_.at(axis);
        auto out = module.forward({X}).toTuple();
        auto order = out->elements().at(0).toTensor();
        auto modes = out->elements().at(1).toTensor();
        return {order, modes};
    }
    auto &model = models_.at(axis);
    model->eval();

    return model->forward(X);
}

/*
Summary:
  Infer outputs for all axes from one `[axes x features]` tensor.
Args:
  features: Input tensor with one row per axis.
Returns:
  `std::pair<torch::Tensor, torch::Tensor>` concatenated order and mode outputs.
Side Effects:
  Switches each axis model to evaluation mode.
Raises:
  std::invalid_argument: If input tensor rank or shape is invalid.
Preconditions:
  `features.rows == axes_` and `features.cols == in_features_`.
*/
std::pair<torch::Tensor, torch::Tensor>
MapFitter::inferBatch(const torch::Tensor &features) {
    auto X = features.contiguous().to(torch::kFloat32);
    if (X.dim() != 2) {
        throw std::invalid_argument("inferBatch expects a 2D feature tensor");
    }
    if (X.size(1) != in_features_) {
        std::ostringstream oss;
        oss << "inferBatch expected feature width " << in_features_
            << ", got " << X.size(1);
        throw std::invalid_argument(oss.str());
    }
    if (X.size(0) != axes_) {
        std::ostringstream oss;
        oss << "inferBatch expected " << axes_ << " rows, got " << X.size(0);
        throw std::invalid_argument(oss.str());
    }

    std::vector<torch::Tensor> orders;
    std::vector<torch::Tensor> modes;
    orders.reserve(axes_);
    modes.reserve(axes_);

    for (int axis = 0; axis < axes_; ++axis) {


        auto row = X[axis].unsqueeze(0);  // [1, in_features]
        if (scripted_) {
            auto &module = scripted_models_.at(axis);
            auto out = module.forward({row}).toTuple();
            orders.push_back(out->elements().at(0).toTensor());
            modes.push_back(out->elements().at(1).toTensor());
        } else {
            auto &model = models_.at(axis);
            model->eval();
            auto out = model->forward(row);
            orders.push_back(out.first);
            modes.push_back(out.second);
        }

    }

    auto order_concat = torch::cat(orders, 0);
    auto modes_concat = torch::cat(modes, 0);
    return {order_concat, modes_concat};
}

/*
Summary:
  Load axis-specific model weights from disk.
Args:
  directory: Path containing `axis{i}_model.pt` files.
Returns:
  None.
Side Effects:
  Replaces in-memory model parameters.
Raises:
  torch::Error / std::exception: If checkpoint loading fails.
Preconditions:
  Checkpoint files exist and match current model architecture.
*/
void MapFitter::load_models(const std::string &directory) {

    scripted_ = false;
    scripted_models_.clear();

    bool scripted_ok = true;
    for (size_t i = 0; i < models_.size(); ++i) {
        std::ostringstream fn;
        fn << directory << "/axis" << i << "_model.pt";
        try {
            auto module = torch::jit::load(fn.str());
            scripted_models_.push_back(std::move(module));
        } catch (const c10::Error&) {
            scripted_ok = false;
            break;
        }
    }

    if (scripted_ok) {
        scripted_ = true;
        return;
    }

    scripted_models_.clear();

    for (size_t i = 0; i < models_.size(); ++i) {
        std::ostringstream fn;
        fn << directory << "/axis" << i << "_model.pt";
        torch::load(models_[i], fn.str());
    }
}

// --- loader ---------------------------------------------------------------

/*
Summary:
  Build a model loader instance and populate it from stored checkpoints.
Args:
  directory: Path containing model files.
  axes: Number of axis models.
  input_features: Number of input features per model.
  hidden: Hidden layer widths.
Returns:
  `std::shared_ptr<MapFitter>` loaded fitter.
Side Effects:
  Reads model files from disk.
Raises:
  torch::Error / std::exception: If loading fails.
Preconditions:
  Saved files are compatible with provided architecture arguments.
*/
std::shared_ptr<MapFitter> ModelLoader::load(const std::string &directory,
                                             int axes,
                                             int input_features,
                                             std::vector<int> hidden) {
    auto fitter = std::make_shared<MapFitter>(axes, input_features, std::move(hidden));
    fitter->load_models(directory);
    return fitter;
}

} // namespace maploader
