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

namespace concord {
namespace utils {

void soft_threshold_matrix(MatrixXd& X, double tau, const MatrixXd& Lambda) {
    for (int i = 0; i < X.rows(); ++i) {
        for (int j = 0; j < X.cols(); ++j) {
            X(i, j) = soft_threshold(X(i, j), tau * Lambda(i, j));
        }
    }
}

bool is_positive_definite(const MatrixXd& X) {
    // Check if all diagonal elements are positive
    for (int i = 0; i < X.rows(); ++i) {
        if (X(i, i) <= 0.0) {
            return false;
        }
    }
    
    // For a more thorough check, we could compute eigenvalues,
    // but diagonal check is sufficient for our iterative algorithms
    // that maintain positive definiteness by construction
    return true;
}

double compute_objective(const MatrixXd& X, const MatrixXd& S, const MatrixXd& Lambda) {
    // Note: This function computes the ISTA objective: -log det(X) + 0.5*tr(X*S*X) + lambda*||X||_1
    // For coordinate descent objective (tr(S*X) instead), use compute_objective_cd
    
    // Compute log determinant
    double log_det = 0.0;
    for (int i = 0; i < X.rows(); ++i) {
        if (X(i, i) <= 0.0) {
            return std::numeric_limits<double>::infinity();
        }
        log_det += std::log(X(i, i));
    }
    
    // Compute 0.5 * trace(X * S * X)
    MatrixXd XSX = X * S * X;
    double trace_xsx = 0.5 * XSX.trace();
    
    // Compute L1 penalty
    double l1_penalty = 0.0;
    for (int i = 0; i < X.rows(); ++i) {
        for (int j = 0; j < X.cols(); ++j) {
            l1_penalty += Lambda(i, j) * std::abs(X(i, j));
        }
    }
    
    return -log_det + trace_xsx + l1_penalty;
}

double compute_objective_cd(const MatrixXd& X, const MatrixXd& S, const MatrixXd& Lambda) {
    // Coordinate descent objective: -log det(X) + tr(S*X) + lambda*||X||_1
    
    // Compute log determinant
    double log_det = 0.0;
    for (int i = 0; i < X.rows(); ++i) {
        if (X(i, i) <= 0.0) {
            return std::numeric_limits<double>::infinity();
        }
        log_det += std::log(X(i, i));
    }
    
    // Compute trace(S * X)
    double trace_sx = (S * X).trace();
    
    // Compute L1 penalty
    double l1_penalty = 0.0;
    for (int i = 0; i < X.rows(); ++i) {
        for (int j = 0; j < X.cols(); ++j) {
            l1_penalty += Lambda(i, j) * std::abs(X(i, j));
        }
    }
    
    return -log_det + trace_sx + l1_penalty;
}

/**
 * @brief Compute gradient of smooth part h1(X) = -log det(X) + 0.5 * tr(X * S * X)
 */
MatrixXd compute_h1_gradient(const MatrixXd& X, const MatrixXd& S) {
    MatrixXd grad = MatrixXd::Zero(X.rows(), X.cols());
    
    // Gradient of -log det(X): -X^{-1} (diagonal terms only for our case)
    for (int i = 0; i < X.rows(); ++i) {
        grad(i, i) = -1.0 / X(i, i);
    }
    
    // Gradient of 0.5 * tr(X * S * X): 0.5 * (S * X + X * S)
    MatrixXd SX = S * X;
    grad += 0.5 * (SX + SX.transpose());
    
    return grad;
}

/**
 * @brief Compute h1 function value
 */
double compute_h1(const MatrixXd& X, const MatrixXd& S) {
    // h1(X) = -log det(X) + 0.5 * tr(X * S * X)
    
    double log_det = 0.0;
    for (int i = 0; i < X.rows(); ++i) {
        if (X(i, i) <= 0.0) {
            return std::numeric_limits<double>::infinity();
        }
        log_det += std::log(X(i, i));
    }
    
    double trace_xsx = 0.5 * (X * S * X).trace();
    
    return -log_det + trace_xsx;
}

/**
 * @brief Compute subgradient of h2(X) = ||X||_1
 */
MatrixXd compute_h2_subgradient(const MatrixXd& X, const MatrixXd& Lambda, const MatrixXd* grad_h1) {
    MatrixXd subgrad = MatrixXd::Zero(X.rows(), X.cols());
    
    for (int i = 0; i < X.rows(); ++i) {
        for (int j = 0; j < X.cols(); ++j) {
            if (X(i, j) != 0.0) {
                // Subgradient when X_ij != 0: lambda * sign(X_ij)
                subgrad(i, j) = Lambda(i, j) * ((X(i, j) > 0) ? 1.0 : -1.0);
            } else if (grad_h1 != nullptr) {
                // Subgradient when X_ij = 0: choose value in [-lambda, lambda] to minimize total subgradient
                double g = (*grad_h1)(i, j);
                double lambda_val = Lambda(i, j);
                if (std::abs(g) <= lambda_val) {
                    subgrad(i, j) = g;
                } else {
                    subgrad(i, j) = lambda_val * ((g > 0) ? 1.0 : -1.0);
                }
            }
        }
    }
    
    return subgrad;
}

/**
 * @brief Check convergence based on subgradient norm
 */
bool check_subgradient_convergence(const MatrixXd& X, const MatrixXd& S, const MatrixXd& Lambda, double tolerance) {
    MatrixXd grad_h1 = compute_h1_gradient(X, S);
    MatrixXd subgrad_h2 = compute_h2_subgradient(X, Lambda, &grad_h1);
    MatrixXd total_subgrad = grad_h1 + subgrad_h2;
    
    double subgrad_norm = total_subgrad.norm();
    double x_norm = X.norm();
    
    return (subgrad_norm / std::max(x_norm, 1.0)) < tolerance;
}

/**
 * @brief Vectorized element access for symmetric matrices
 */
void extract_lower_triangular(const MatrixXd& X, VectorXd& vec) {
    int p = X.rows();
    int idx = 0;
    for (int i = 0; i < p; ++i) {
        for (int j = i; j < p; ++j) {
            vec(idx++) = X(i, j);
        }
    }
}

void set_from_lower_triangular(MatrixXd& X, const VectorXd& vec) {
    int p = X.rows();
    int idx = 0;
    for (int i = 0; i < p; ++i) {
        for (int j = i; j < p; ++j) {
            X(i, j) = vec(idx);
            X(j, i) = vec(idx);  // Ensure symmetry
            idx++;
        }
    }
}



} // namespace utils
} // namespace concord 