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

namespace concord {

// Forward declarations for internal utility functions
namespace {
    double sign(double val) {
        return (val > 0.0) ? 1.0 : ((val < 0.0) ? -1.0 : 0.0);
    }
    
    void soft_threshold_matrix_ista(MatrixXd& X, double tau, const MatrixXd& Lambda) {
        for (int i = 0; i < X.rows(); ++i) {
            for (int j = 0; j < X.cols(); ++j) {
                double threshold = tau * Lambda(i, j);
                double abs_x = std::abs(X(i, j));
                if (abs_x > threshold) {
                    X(i, j) = sign(X(i, j)) * (abs_x - threshold);
                } else {
                    X(i, j) = 0.0;
                }
            }
        }
    }
}

/**
 * @brief ISTA (Iterative Shrinkage-Thresholding Algorithm) implementation
 * Matrix-based proximal gradient method with backtracking line search
 */
ConcordResult ConcordSolver::solve_ista(const MatrixXd& S, const MatrixXd& Lambda) {
    ConcordResult result;
    result.algorithm_used = Algorithm::ISTA;
    
    int p = S.rows();
    
    // Initialize matrices
    SparseMatrix<double> X(p, p), X_next(p, p);
    SparseMatrix<double> Step(p, p);
    MatrixXd W(p, p), W_next(p, p);
    MatrixXd grad_h1(p, p);
    MatrixXd subgrad_h2(p, p);
    MatrixXd subgrad_total(p, p);
    MatrixXd tmp(p, p);
    
    // Algorithm parameters (from original pyccista)
    double tau = 1.0;
    double c = 0.5;  // From original pyccista
    
    // Initialize X as identity matrix
    X.setIdentity();
    W = S * MatrixXd(X);
    
    // Compute initial gradient and objective
    grad_h1 = MatrixXd::Zero(p, p);
    for (int i = 0; i < p; ++i) {
        grad_h1(i, i) = -1.0 / X.coeff(i, i);
    }
    grad_h1 += 0.5 * (W + W.transpose());
    
    double h1 = 0.0;
    for (int i = 0; i < p; ++i) {
        h1 -= std::log(X.coeff(i, i));
    }
    h1 += 0.5 * MatrixXd(X).cwiseProduct(W).sum();
    
    bool converged = false;
    int outer_iteration = 0;
    
    // Store initial objective if requested
    if (options_.save_iterates) {
        double initial_objective = h1;
        for (int i = 0; i < p; ++i) {
            for (int j = 0; j < p; ++j) {
                initial_objective += Lambda(i, j) * std::abs(X.coeff(i, j));
            }
        }
        result.objective_values.push_back(initial_objective);
        result.iterates.push_back(MatrixXd(X));
    }
    
    while (outer_iteration < options_.max_iterations && !converged) {
        tau = 1.0;  // Reset step size (from original pyccista)
        int inner_iteration = 0;
        
        // Backtracking line search
        while (true) {
            // Proximal gradient step
            tmp = MatrixXd(X) - tau * grad_h1;
            soft_threshold_matrix_ista(tmp, tau, Lambda);
            X_next = tmp.sparseView();
            
            // Check if diagonal elements are positive
            bool positive_diagonal = true;
            for (int i = 0; i < p; ++i) {
                if (tmp(i, i) <= 0.0) {
                    positive_diagonal = false;
                    break;
                }
            }
            
            if (positive_diagonal) {
                // Compute new objective value
                Step = X_next - X;
                W_next = S * MatrixXd(X_next);
                
                double h1_next = 0.0;
                for (int i = 0; i < p; ++i) {
                    h1_next -= std::log(X_next.coeff(i, i));
                }
                h1_next += 0.5 * MatrixXd(X_next).cwiseProduct(W_next).sum();
                
                // Quadratic approximation for sufficient decrease condition
                double Q = h1 + MatrixXd(Step).cwiseProduct(grad_h1).sum() + 
                          (0.5 / tau) * MatrixXd(Step).squaredNorm();
                
                if (h1_next <= Q) {
                    break;  // Sufficient decrease achieved
                }
            }
            
            // Backtrack
            tau *= c;
            inner_iteration++;
        }
        
        // Update gradient
        grad_h1 = MatrixXd::Zero(p, p);
        for (int i = 0; i < p; ++i) {
            grad_h1(i, i) = -1.0 / X_next.coeff(i, i);
        }
        grad_h1 += 0.5 * (W_next + W_next.transpose());
        
        // Compute subgradient for convergence check
        subgrad_h2 = MatrixXd::Zero(p, p);
        
        // Subgradient when X_ij is zero: choose value in [-lambda, lambda] to minimize total subgradient
        for (int i = 0; i < p; ++i) {
            for (int j = 0; j < p; ++j) {
                if (X_next.coeff(i, j) == 0.0) {
                    double g = grad_h1(i, j);
                    double lambda_val = Lambda(i, j);
                    if (std::abs(g) <= lambda_val) {
                        subgrad_h2(i, j) = g;
                    } else {
                        subgrad_h2(i, j) = lambda_val * sign(g);
                    }
                } else {
                    // Subgradient when X_ij != 0: lambda * sign(X_ij)
                    subgrad_h2(i, j) = Lambda(i, j) * sign(X_next.coeff(i, j));
                }
            }
        }
        
        subgrad_total = grad_h1 + subgrad_h2;
        double delta_subgrad = subgrad_total.norm() / std::max(MatrixXd(X_next).norm(), 1.0);
        
        // Compute full objective value
        double h1_next = 0.0;
        for (int i = 0; i < p; ++i) {
            h1_next -= std::log(X_next.coeff(i, i));
        }
        h1_next += 0.5 * MatrixXd(X_next).cwiseProduct(W_next).sum();
        
        double full_objective = h1_next;
        for (int i = 0; i < p; ++i) {
            for (int j = 0; j < p; ++j) {
                full_objective += Lambda(i, j) * std::abs(X_next.coeff(i, j));
            }
        }
        
        // Store iteration information
        result.inner_iterations.push_back(inner_iteration);
        result.convergence_metrics.push_back(delta_subgrad);
        if (options_.save_iterates) {
            result.objective_values.push_back(full_objective);
            result.iterates.push_back(MatrixXd(X_next));
        }
        
        // Check convergence based on subgradient norm
        if (delta_subgrad < options_.tolerance) {
            converged = true;
            result.convergence_reason = "Subgradient norm below tolerance";
        }
        
        // Update for next iteration
        h1 = h1_next;
        X = X_next;
        W = W_next;
        
        outer_iteration++;
        
        if (options_.verbose && outer_iteration % 10 == 0) {
            std::cout << "ISTA iteration " << outer_iteration 
                      << ", subgrad_norm = " << delta_subgrad 
                      << ", inner_iters = " << inner_iteration << std::endl;
        }
    }
    
    // Finalize result
    result.omega = X;
    result.total_iterations = outer_iteration;
    result.converged = converged;
    result.final_tolerance = result.convergence_metrics.back();
    
    if (!converged) {
        result.convergence_reason = "Maximum iterations reached";
    }
    
    return result;
}

} // namespace concord 