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

namespace concord {

/**
 * @brief Symmetric Lasso implementation (symlassoC)
 * Uses auxiliary variables for efficient diagonal updates
 */
ConcordResult ConcordSolver::solve_symlasso(const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::SYMLASSO;
    
    int p = S.rows();
    // int tp = p * (p + 1) / 2;  // Unused variable
    
    // Initialize matrices and auxiliary variables
    MatrixXd omega = MatrixXd::Zero(p, p);
    MatrixXd Z = MatrixXd::Zero(p, p);
    VectorXd q = VectorXd::Ones(p);
    VectorXd sig = VectorXd::Ones(p);
    
    // Use scalar lambda (symlasso uses single lambda value)
    double lambda_scalar = Lambda(0, 0);  // Extract scalar from matrix
    
    bool converged = false;
    int iteration = 0;
    double max_diff = 1.0;
    
    // Store initial objective if requested
    if (options_.save_iterates) {
        // Construct full omega matrix for objective computation
        MatrixXd omega_full = omega;
        for (int i = 0; i < p; ++i) {
            omega_full(i, i) = 1.0 / sig(i);
        }
        result.objective_values.push_back(utils::compute_objective(omega_full, S, Lambda));
        result.iterates.push_back(omega_full);
    }
    
    while (iteration < options_.max_iterations && !converged) {
        max_diff = 0.0;
        
        // Update off-diagonal elements
        for (int i = 0; i < p - 1; ++i) {
            for (int j = i + 1; j < p; ++j) {
                // Compute sij and sji using auxiliary variables
                double sij = S(i, j) + Z(i, j) * sig(j) - omega(i, j) * sig(j) * S(i, i);
                double sji = S(j, i) + Z(j, i) * sig(i) - omega(j, i) * sig(i) * S(j, j);
                
                double delta_ij = omega(i, j);
                
                // Apply soft thresholding
                double denominator = S(j, j) * sig(i) + S(i, i) * sig(j);
                double new_val = utils::soft_threshold(-(sij + sji), lambda_scalar) / denominator;
                
                omega(i, j) = new_val;
                omega(j, i) = new_val;
                
                delta_ij = new_val - delta_ij;  // delta = omega_new - omega_old
                
                if (delta_ij != 0.0) {
                    // Update auxiliary variables q and Z
                    q(j) += 2 * Z(i, j) * delta_ij + S(i, i) * delta_ij * delta_ij;
                    q(i) += 2 * Z(j, i) * delta_ij + S(j, j) * delta_ij * delta_ij;
                    
                    // Update Z matrix
                    for (int k = 0; k < p; ++k) {
                        Z(k, j) += S(i, k) * delta_ij;
                        Z(k, i) += S(j, k) * delta_ij;
                    }
                }
                
                double diff = std::abs(delta_ij);
                max_diff = std::max(max_diff, diff);
            }
        }
        
        // Update diagonal scaling factors
        for (int i = 0; i < p; ++i) {
            double old_sig = sig(i);
            sig(i) = (-1 + std::sqrt(1 + 4 * q(i) * S(i, i))) / (2 * q(i));
            
            double diff = std::abs(sig(i) - old_sig);
            max_diff = std::max(max_diff, diff);
        }
        
        // Store iteration information
        result.convergence_metrics.push_back(max_diff);
        if (options_.save_iterates) {
            // Construct full omega matrix for objective computation
            MatrixXd omega_full = omega;
            for (int i = 0; i < p; ++i) {
                omega_full(i, i) = 1.0 / sig(i);
            }
            result.objective_values.push_back(utils::compute_objective(omega_full, S, Lambda));
            result.iterates.push_back(omega_full);
        }
        
        // Check convergence
        if (max_diff < options_.tolerance) {
            converged = true;
            result.convergence_reason = "Maximum element change below tolerance";
        }
        
        iteration++;
        
        if (options_.verbose && iteration % 10 == 0) {
            std::cout << "Symlasso iteration " << iteration << ", max_diff = " << max_diff << std::endl;
        }
    }
    
    // Set final diagonal elements
    for (int i = 0; i < p; ++i) {
        omega(i, i) = 1.0 / sig(i);
    }
    
    // Finalize result
    result.omega = omega.sparseView();
    result.total_iterations = iteration;
    result.converged = converged;
    result.final_tolerance = max_diff;
    
    if (!converged) {
        result.convergence_reason = "Maximum iterations reached";
    }
    
    return result;
}

/**
 * @brief Symmetric Lasso with iterate saving (symlassoC2)
 * Variant that saves detailed iteration history
 */
ConcordResult ConcordSolver::solve_symlasso_iterates(const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::SYMLASSO_ITERATES;
    
    int p = S.rows();
    int tp = p * (p + 1) / 2;
    
    // Initialize matrices and auxiliary variables
    MatrixXd omega = MatrixXd::Zero(p, p);
    MatrixXd Z = MatrixXd::Zero(p, p);
    VectorXd q = VectorXd::Ones(p);
    VectorXd sig = VectorXd::Ones(p);
    
    // Use scalar lambda
    double lambda_scalar = Lambda(0, 0);
    
    bool converged = false;
    int iteration = 1;  // Start from 1 as in original implementation
    double max_diff = 1.0;
    
    // Pre-allocate iterate storage
    std::vector<VectorXd> iterate_history;
    VectorXd current_iterate(tp);
    VectorXd previous_iterate = VectorXd::Zero(tp);
    
    while (iteration <= options_.max_iterations - 1 && !converged) {
        // Update off-diagonal elements
        for (int i = 0; i < p - 1; ++i) {
            for (int j = i + 1; j < p; ++j) {
                double sij = S(i, j) + Z(i, j) * sig(j) - omega(i, j) * sig(j) * S(i, i);
                double sji = S(j, i) + Z(j, i) * sig(i) - omega(j, i) * sig(i) * S(j, j);
                
                double delta_ij = omega(i, j);
                
                double denominator = S(j, j) * sig(i) + S(i, i) * sig(j);
                double new_val = utils::soft_threshold(-(sij + sji), lambda_scalar) / denominator;
                
                omega(i, j) = new_val;
                omega(j, i) = new_val;
                
                delta_ij = new_val - delta_ij;
                
                if (delta_ij != 0.0) {
                    q(j) += 2 * Z(i, j) * delta_ij + S(i, i) * delta_ij * delta_ij;
                    q(i) += 2 * Z(j, i) * delta_ij + S(j, j) * delta_ij * delta_ij;
                    
                    for (int k = 0; k < p; ++k) {
                        Z(k, j) += S(i, k) * delta_ij;
                        Z(k, i) += S(j, k) * delta_ij;
                    }
                }
            }
        }
        
        // Update diagonal scaling factors
        for (int i = 0; i < p; ++i) {
            sig(i) = (-1 + std::sqrt(1 + 4 * q(i) * S(i, i))) / (2 * q(i));
        }
        
        // Save iterates and compute convergence
        max_diff = 0.0;
        // int idx = 0;  // Unused variable
        
        // Store diagonal elements first
        for (int i = 0; i < p; ++i) {
            int k = utils::linear_index(i, i, p);
            current_iterate(k) = 1.0 / sig(i);
        }
        
        // Store off-diagonal elements
        for (int i = 0; i < p; ++i) {
            for (int j = i; j < p; ++j) {
                int k = utils::linear_index(i, j, p);
                current_iterate(k) = omega(i, j);
                
                if (iteration > 1) {
                    double diff = std::abs(current_iterate(k) - previous_iterate(k));
                    max_diff = std::max(max_diff, diff);
                }
            }
        }
        
        // Store iteration information
        result.convergence_metrics.push_back(max_diff);
        if (options_.save_iterates) {
            iterate_history.push_back(current_iterate);
        }
        
        // Check convergence
        if (iteration > 1 && max_diff < options_.tolerance) {
            converged = true;
            result.convergence_reason = "Maximum element change below tolerance";
        }
        
        previous_iterate = current_iterate;
        iteration++;
        
        if (options_.verbose && iteration % 10 == 0) {
            std::cout << "Symlasso iterates " << iteration << ", max_diff = " << max_diff << std::endl;
        }
    }
    
    // Set final diagonal elements
    for (int i = 0; i < p; ++i) {
        omega(i, i) = 1.0 / sig(i);
    }
    
    // Convert iterate history to matrix format if requested
    if (options_.save_iterates && !iterate_history.empty()) {
        result.iterates.reserve(iterate_history.size());
        for (const auto& iter : iterate_history) {
            MatrixXd iter_matrix = MatrixXd::Zero(p, p);
            utils::set_from_lower_triangular(iter_matrix, iter);
            result.iterates.push_back(iter_matrix);
        }
    }
    
    // Finalize result
    result.omega = omega.sparseView();
    result.total_iterations = iteration - 1;
    result.converged = converged;
    result.final_tolerance = max_diff;
    
    if (!converged) {
        result.convergence_reason = "Maximum iterations reached";
    }
    
    return result;
}

} // namespace concord 