#include "concord_unified.hpp"
#include <cmath>
#include <algorithm>
#include <iostream>
#include <stdexcept>

namespace concord {

ConcordSolver::ConcordSolver(const ConcordOptions& options) : options_(options) {
}

Algorithm ConcordSolver::select_algorithm(const MatrixXd& S, const MatrixXd* data_matrix) {
    if (options_.algorithm != Algorithm::AUTO) {
        return options_.algorithm;
    }
    
    int p = S.rows();
    int n = (data_matrix != nullptr) ? data_matrix->rows() : p + 1;  // Assume n > p if no data provided
    
    // Use original R package logic only: if (n <= p) use concordC2, else concordC
    if (n <= p) {
        // High-dimensional case: use coordinate descent with residual updates (concordC2)
        if (options_.verbose) {
            std::cout << "Auto-selecting COORDINATE_HIGH_DIM (n=" << n << " <= p=" << p << ")" << std::endl;
        }
        return Algorithm::COORDINATE_HIGH_DIM;
    } else {
        // Standard case: use coordinate descent (concordC)
        if (options_.verbose) {
            std::cout << "Auto-selecting COORDINATE_DESCENT (n=" << n << " > p=" << p << ")" << std::endl;
        }
        return Algorithm::COORDINATE_DESCENT;
    }
}

MatrixXd ConcordSolver::create_lambda_matrix(double lambda_scalar, int p) {
    MatrixXd Lambda = MatrixXd::Constant(p, p, lambda_scalar);
    
    if (!options_.penalize_diagonal) {
        Lambda.diagonal().setZero();
    }
    
    return Lambda;
}

ConcordResult ConcordSolver::solve(const MatrixXd& S, const MatrixXd& Lambda) {
    // Validate inputs
    if (S.rows() != S.cols()) {
        throw std::invalid_argument("Sample covariance matrix S must be square");
    }
    
    if (Lambda.rows() != S.rows() || Lambda.cols() != S.cols()) {
        throw std::invalid_argument("Lambda matrix must have same dimensions as S");
    }
    
    if (!S.isApprox(S.transpose())) {
        throw std::invalid_argument("Sample covariance matrix S must be symmetric");
    }
    
    if (!Lambda.isApprox(Lambda.transpose())) {
        throw std::invalid_argument("Lambda matrix must be symmetric");
    }
    
    // Select algorithm
    Algorithm selected_algorithm = select_algorithm(S, nullptr);
    
    // Dispatch to appropriate solver
    switch (selected_algorithm) {
        case Algorithm::COORDINATE_DESCENT:
            return solve_coordinate_descent(S, Lambda);
            
        case Algorithm::COORDINATE_BLAS:
            return solve_coordinate_blas(S, Lambda);
            
        case Algorithm::ISTA:
            return solve_ista(S, Lambda);
            
        case Algorithm::SYMLASSO:
            return solve_symlasso(S, Lambda);
            
        case Algorithm::SYMLASSO_ITERATES:
            return solve_symlasso_iterates(S, Lambda);
            
        case Algorithm::COORDINATE_HIGH_DIM:
            throw std::invalid_argument("COORDINATE_HIGH_DIM requires raw data matrix. Use solve_from_data() instead.");
            
        default:
            throw std::invalid_argument("Unknown algorithm selected");
    }
}

ConcordResult ConcordSolver::solve(const MatrixXd& S, double lambda_scalar) {
    MatrixXd Lambda = create_lambda_matrix(lambda_scalar, S.rows());
    return solve(S, Lambda);
}

ConcordResult ConcordSolver::solve_from_data(const MatrixXd& data, double lambda_scalar) {
    int n = data.rows();
    int p = data.cols();
    
    // Compute sample covariance matrix
    MatrixXd S;
    if (n <= p) {
        // High-dimensional case: use diagonal covariance
        S = MatrixXd::Zero(p, p);
        for (int i = 0; i < p; ++i) {
            S(i, i) = data.col(i).squaredNorm() / n;
        }
    } else {
        // Standard case: full covariance matrix
        MatrixXd centered_data = data.rowwise() - data.colwise().mean();
        S = (centered_data.transpose() * centered_data) / n;
    }
    
    MatrixXd Lambda = create_lambda_matrix(lambda_scalar, p);
    
    // Select algorithm with data information
    Algorithm selected_algorithm = select_algorithm(S, &data);
    
    // Dispatch to appropriate solver
    switch (selected_algorithm) {
        case Algorithm::COORDINATE_HIGH_DIM:
            return solve_coordinate_high_dim(data, S, Lambda);
            
        case Algorithm::COORDINATE_DESCENT:
            return solve_coordinate_descent(S, Lambda);
            
        case Algorithm::COORDINATE_BLAS:
            return solve_coordinate_blas(S, Lambda);
            
        case Algorithm::ISTA:
            return solve_ista(S, Lambda);
            
        case Algorithm::SYMLASSO:
            return solve_symlasso(S, Lambda);
            
        case Algorithm::SYMLASSO_ITERATES:
            return solve_symlasso_iterates(S, Lambda);
            
        default:
            throw std::invalid_argument("Unknown algorithm selected");
    }
}

// Convenience functions
ConcordResult concord(const MatrixXd& S, double lambda_scalar, const ConcordOptions& options) {
    ConcordSolver solver(options);
    return solver.solve(S, lambda_scalar);
}

ConcordResult concord_from_data(const MatrixXd& data, double lambda_scalar, const ConcordOptions& options) {
    ConcordSolver solver(options);
    return solver.solve_from_data(data, lambda_scalar);
}

} // namespace concord 