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

namespace concord {

/**
 * @brief Standard coordinate descent implementation (concordC)
 * For n > p case with full covariance matrix
 */
ConcordResult ConcordSolver::solve_coordinate_descent(const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::COORDINATE_DESCENT;
    
    int p = S.rows();
    int tp = p * (p + 1) / 2;  // Lower triangular storage size
    
    // Initialize omega as identity matrix
    MatrixXd omega = MatrixXd::Identity(p, p);
    MatrixXd omega_old = MatrixXd::Zero(p, p);
    
    // Convert lambda matrix to vector for compatibility with original indexing
    VectorXd lambda_vec(tp);
    int idx = 0;
    for (int i = 0; i < p; ++i) {
        for (int j = i; j < p; ++j) {
            lambda_vec(idx++) = Lambda(i, j);
        }
    }
    
    // Scale lambda by factor (from original R interface)
    lambda_vec *= options_.lambda_scale;
    
    bool converged = false;
    int iteration = 0;
    double max_diff = 1.0;
    
    // Store initial objective if requested
    if (options_.save_iterates) {
        result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
        result.iterates.push_back(omega);
    }
    
    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) {
                double s1 = 0.0, s2 = 0.0;
                
                // Compute s1 = sum_k omega(i,k) * S(j,k) - omega(i,j) * S(j,j)
                for (int k = 0; k < p; ++k) {
                    s1 += omega(i, k) * S(j, k);
                    s2 += omega(k, j) * S(i, k);
                }
                s1 -= omega(i, j) * S(j, j);
                s2 -= omega(i, j) * S(i, i);
                
                double s_sum = S(i, i) + S(j, j);
                int linear_idx = utils::linear_index(i, j, p);
                
                // Apply soft thresholding
                double new_val = utils::soft_threshold(-(s1 + s2) / s_sum, 
                                                     lambda_vec(linear_idx) / s_sum);
                
                // Update both (i,j) and (j,i) for symmetry
                double diff = std::abs(new_val - omega(i, j));
                max_diff = std::max(max_diff, diff);
                
                omega(i, j) = new_val;
                omega(j, i) = new_val;
            }
        }
        
        // Update diagonal elements
        for (int i = 0; i < p; ++i) {
            double s1 = 0.0;
            for (int k = 0; k < p; ++k) {
                s1 += omega(i, k) * S(i, k);
            }
            s1 -= omega(i, i) * S(i, i);
            
            // Quadratic formula solution to maintain positive definiteness
            double new_val = (-s1 + std::sqrt(s1 * s1 + 4 * S(i, i))) / (2 * S(i, i));
            
            double diff = std::abs(new_val - omega(i, i));
            max_diff = std::max(max_diff, diff);
            
            omega(i, i) = new_val;
        }
        
        // Store iteration information
        result.convergence_metrics.push_back(max_diff);
        if (options_.save_iterates) {
            result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
            result.iterates.push_back(omega);
        }
        
        // 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 << "Iteration " << iteration << ", max_diff = " << max_diff << std::endl;
        }
    }
    
    // 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 High-dimensional coordinate descent implementation (concordC2)
 * For n <= p case with residual updates
 */
ConcordResult ConcordSolver::solve_coordinate_high_dim(const MatrixXd& data, const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::COORDINATE_HIGH_DIM;
    
    int n = data.rows();
    int p = data.cols();
    int tp = p * (p + 1) / 2;
    
    // Initialize omega and residuals
    MatrixXd omega = MatrixXd::Identity(p, p);
    MatrixXd omega_old = MatrixXd::Zero(p, p);
    MatrixXd residuals = data;  // Copy of data matrix
    
    // Convert lambda matrix to vector
    VectorXd lambda_vec(tp);
    int idx = 0;
    for (int i = 0; i < p; ++i) {
        for (int j = i; j < p; ++j) {
            lambda_vec(idx++) = Lambda(i, j);
        }
    }
    lambda_vec *= options_.lambda_scale;
    
    bool converged = false;
    int iteration = 0;
    double max_diff = 1.0;
    
    // Store initial objective if requested
    if (options_.save_iterates) {
        result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
        result.iterates.push_back(omega);
    }
    
    while (iteration < options_.max_iterations && !converged) {
        max_diff = 0.0;
        
        // Update off-diagonal elements with residual updates
        for (int i = 0; i < p - 1; ++i) {
            for (int j = i + 1; j < p; ++j) {
                // Compute s1 and s2 using residuals and dot products
                double s1 = -omega(i, j) * S(j, j) + 
                           omega(i, i) * residuals.col(i).dot(data.col(j)) / n;
                double s2 = -omega(j, i) * S(i, i) + 
                           omega(j, j) * residuals.col(j).dot(data.col(i)) / n;
                
                double s_sum = S(i, i) + S(j, j);
                int linear_idx = utils::linear_index(i, j, p);
                
                double new_val = utils::soft_threshold(-(s1 + s2) / s_sum, 
                                                     lambda_vec(linear_idx) / s_sum);
                
                double tmp_diff = new_val - omega(i, j);
                
                if (tmp_diff != 0.0) {
                    // Update residuals efficiently
                    double scale1 = tmp_diff / omega(i, i);
                    double scale2 = tmp_diff / omega(j, j);
                    
                    residuals.col(i) += scale1 * data.col(j);
                    residuals.col(j) += scale2 * data.col(i);
                    
                    max_diff = std::max(max_diff, std::abs(tmp_diff));
                }
                
                omega(i, j) = new_val;
                omega(j, i) = new_val;
            }
        }
        
        // Update diagonal elements
        for (int i = 0; i < p; ++i) {
            double s1 = -omega(i, i) * S(i, i) + 
                       omega(i, i) * residuals.col(i).dot(data.col(i)) / n;
            
            double new_val = (-s1 + std::sqrt(s1 * s1 + 4 * S(i, i))) / (2 * S(i, i));
            
            // Update residuals for diagonal change
            double scale = omega(i, i) / new_val;
            residuals.col(i) *= scale;
            residuals.col(i) += (1 - scale) * data.col(i);
            
            double diff = std::abs(new_val - omega(i, i));
            max_diff = std::max(max_diff, diff);
            
            omega(i, i) = new_val;
        }
        
        // Store iteration information
        result.convergence_metrics.push_back(max_diff);
        if (options_.save_iterates) {
            result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
            result.iterates.push_back(omega);
        }
        
        // 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 << "High-dim iteration " << iteration << ", max_diff = " << max_diff << std::endl;
        }
    }
    
    // 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 BLAS-optimized coordinate descent implementation (concordC3)
 * Uses Eigen's optimized matrix operations
 */
ConcordResult ConcordSolver::solve_coordinate_blas(const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::COORDINATE_BLAS;
    
    int p = S.rows();
    int tp = p * (p + 1) / 2;
    
    // Initialize omega as identity matrix
    MatrixXd omega = MatrixXd::Identity(p, p);
    MatrixXd omega_old = MatrixXd::Zero(p, p);
    
    // Convert lambda matrix to vector
    VectorXd lambda_vec(tp);
    int idx = 0;
    for (int i = 0; i < p; ++i) {
        for (int j = i; j < p; ++j) {
            lambda_vec(idx++) = Lambda(i, j);
        }
    }
    lambda_vec *= options_.lambda_scale;
    
    bool converged = false;
    int iteration = 0;
    double max_diff = 1.0;
    
    // Store initial objective if requested
    if (options_.save_iterates) {
        result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
        result.iterates.push_back(omega);
    }
    
    while (iteration < options_.max_iterations && !converged) {
        max_diff = 0.0;
        
        // Update off-diagonal elements using vectorized operations
        for (int i = 0; i < p - 1; ++i) {
            for (int j = i + 1; j < p; ++j) {
                // Use Eigen's dot product for efficiency
                double s1 = omega.row(i).dot(S.col(j)) - omega(i, j) * S(j, j);
                double s2 = omega.row(j).dot(S.col(i)) - omega(i, j) * S(i, i);
                
                double s_sum = S(i, i) + S(j, j);
                int linear_idx = utils::linear_index(i, j, p);
                
                double new_val = utils::soft_threshold(-(s1 + s2) / s_sum, 
                                                     lambda_vec(linear_idx) / s_sum);
                
                double diff = std::abs(new_val - omega(i, j));
                max_diff = std::max(max_diff, diff);
                
                omega(i, j) = new_val;
                omega(j, i) = new_val;
            }
        }
        
        // Update diagonal elements
        for (int i = 0; i < p; ++i) {
            double s1 = omega.row(i).dot(S.col(i)) - omega(i, i) * S(i, i);
            double new_val = (-s1 + std::sqrt(s1 * s1 + 4 * S(i, i))) / (2 * S(i, i));
            
            double diff = std::abs(new_val - omega(i, i));
            max_diff = std::max(max_diff, diff);
            
            omega(i, i) = new_val;
        }
        
        // Store iteration information
        result.convergence_metrics.push_back(max_diff);
        if (options_.save_iterates) {
            result.objective_values.push_back(utils::compute_objective_cd(omega, S, Lambda));
            result.iterates.push_back(omega);
        }
        
        // 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 << "BLAS iteration " << iteration << ", max_diff = " << max_diff << std::endl;
        }
    }
    
    // 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;
}

} // namespace concord 