/***************************************************************************

  PFunction.cpp

  Author: Andreas Suter
  e-mail: andreas.suter@psi.ch

  Modernized implementation using Boost.Spirit X3 with full PMetaData support.

***************************************************************************/

/***************************************************************************
 *   Copyright (C) 2007-2026 by Andreas Suter                              *
 *   andreas.suter@psi.ch                                                  *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <cmath>
#include <iostream>
#include <sstream>

#include <boost/spirit/home/x3.hpp>
#include <boost/variant/apply_visitor.hpp>

#include "PMusr.h"
#include "PFunction.h"

namespace x3 = boost::spirit::x3;

//--------------------------------------------------------------------------
/**
 * \brief Constructor that initializes the function from a parsed AST.
 *
 * This constructor stores the AST generated by the Boost.Spirit X3 parser,
 * extracts the function number, and prepares the function for evaluation.
 * The X3 AST is directly suitable for evaluation using the visitor pattern,
 * eliminating the need for conversion to an intermediate tree structure.
 *
 * \param assignment AST from parsing a function expression (FUN# = expression)
 */
PFunction::PFunction(const musrfit::ast::assignment& assignment)
{
  fValid  = true;
  fFuncNo = assignment.func_number;
  fAst    = assignment.rhs;

  // Generate human-readable string representation
  fFuncString = "FUN";
  fFuncString += fFuncNo;
  fFuncString += " = ";
  fFuncString += GenerateString(fAst);
}

//--------------------------------------------------------------------------
/**
 * \brief Destructor that releases resources.
 *
 * Cleans up parameter and map vectors.
 */
PFunction::~PFunction()
{
  fParam.clear();
  fMap.clear();
}

//--------------------------------------------------------------------------
/**
 * \brief Validates that all parameter and map references are within valid ranges.
 *
 * Recursively traverses the AST to find all parameter (PAR#) and map (MAP#)
 * references and checks that they are within the specified bounds.
 *
 * \param mapSize Number of available map entries
 * \param paramSize Number of available fit parameters
 * \return true if all references are valid, false otherwise
 */
Bool_t PFunction::CheckMapAndParamRange(UInt_t mapSize, UInt_t paramSize)
{
  Bool_t valid = true;

  // Check first operand
  if (!FindAndCheckMapAndParamRange(fAst.first, mapSize, paramSize)) {
    valid = false;
  }

  // Check remaining operations
  for (const auto& op : fAst.rest) {
    if (!FindAndCheckMapAndParamRange(op.operand_, mapSize, paramSize)) {
      valid = false;
    }
  }

  return valid;
}

//--------------------------------------------------------------------------
/**
 * \brief Recursively validates parameter and map references in an operand.
 *
 * Helper visitor for CheckMapAndParamRange that handles different operand types.
 *
 * \param operand Current operand being checked
 * \param mapSize Number of available map entries
 * \param paramSize Number of available fit parameters
 * \return true if all references are valid, false otherwise
 */
Bool_t PFunction::FindAndCheckMapAndParamRange(const musrfit::ast::operand& operand,
                                                 UInt_t mapSize, UInt_t paramSize)
{
  struct CheckVisitor : public boost::static_visitor<Bool_t>
  {
    UInt_t fMapSize;
    UInt_t fParamSize;
    PFunction* fFunc;

    CheckVisitor(UInt_t mapSize, UInt_t paramSize, PFunction* func)
      : fMapSize(mapSize), fParamSize(paramSize), fFunc(func) {}

    Bool_t operator()(const musrfit::ast::nil&) const { return true; }
    Bool_t operator()(double) const { return true; }
    Bool_t operator()(const musrfit::ast::constant&) const { return true; }

    Bool_t operator()(const musrfit::ast::parameter& p) const {
      if (p.number < 1 || static_cast<UInt_t>(p.number) > fParamSize) {
        std::cerr << "**ERROR** Parameter PAR" << p.number << " out of range (1-" << fParamSize << ")" << std::endl;
        return false;
      }
      return true;
    }

    Bool_t operator()(const musrfit::ast::map_ref& m) const {
      if (m.number < 1 || static_cast<UInt_t>(m.number) > fMapSize) {
        std::cerr << "**ERROR** Map MAP" << m.number << " out of range (1-" << fMapSize << ")" << std::endl;
        return false;
      }
      return true;
    }

    Bool_t operator()(const musrfit::ast::function_call& f) const {
      return fFunc->FindAndCheckMapAndParamRange(f.arg.first, fMapSize, fParamSize);
    }

    Bool_t operator()(const musrfit::ast::power_call& p) const {
      Bool_t valid = true;
      if (!fFunc->FindAndCheckMapAndParamRange(p.base.first, fMapSize, fParamSize))
        valid = false;
      if (!fFunc->FindAndCheckMapAndParamRange(p.pow.first, fMapSize, fParamSize))
        valid = false;
      return valid;
    }

    Bool_t operator()(const musrfit::ast::expression& expr) const {
      Bool_t valid = true;
      if (!fFunc->FindAndCheckMapAndParamRange(expr.first, fMapSize, fParamSize))
        valid = false;
      for (const auto& op : expr.rest) {
        if (!fFunc->FindAndCheckMapAndParamRange(op.operand_, fMapSize, fParamSize))
          valid = false;
      }
      return valid;
    }
  };

  CheckVisitor checker(mapSize, paramSize, this);
  return boost::apply_visitor(checker, operand);
}

//--------------------------------------------------------------------------
/**
 * \brief Evaluates the function with given parameters and metadata.
 *
 * Uses the visitor pattern to traverse the AST and compute the numeric result.
 * The evaluation applies operators with proper precedence as encoded in the AST.
 *
 * \param param Vector of fit parameter values
 * \param metaData Metadata containing field, energy, temperature, etc.
 * \return Computed function value, or 0.0 if invalid
 */
Double_t PFunction::Eval(std::vector<Double_t> param, PMetaData metaData)
{
  if (!fValid) return 0.0;

  fParam = param;
  fMetaData = metaData;

  EvalVisitor evaluator(fMap, fParam, fMetaData);

  // Evaluate first operand
  Double_t result = boost::apply_visitor(evaluator, fAst.first);

  // Apply operations
  for (const auto& op : fAst.rest) {
    Double_t rhs = boost::apply_visitor(evaluator, op.operand_);

    switch (op.operator_) {
      case musrfit::ast::op_plus:
        result += rhs;
        break;
      case musrfit::ast::op_minus:
        result -= rhs;
        break;
      case musrfit::ast::op_times:
        result *= rhs;
        break;
      case musrfit::ast::op_divide:
        if (std::abs(rhs) > 1.0e-20)
          result /= rhs;
        else
          result = 0.0;
        break;
    }
  }

  return result;
}

//--------------------------------------------------------------------------
/**
 * \brief Generates a human-readable string representation of an expression.
 *
 * \param expr Expression AST to convert
 * \return Formatted string
 */
TString PFunction::GenerateString(const musrfit::ast::expression& expr)
{
  TString result;

  // Check if we have any high-precedence operations (multiplication/division)
  bool hasHighPrecOps = false;
  for (const auto& op : expr.rest) {
    if (op.operator_ == musrfit::ast::op_times || op.operator_ == musrfit::ast::op_divide) {
      hasHighPrecOps = true;
      break;
    }
  }

  // Generate first operand - check if it's a wrapped expression
  if (const musrfit::ast::expression* inner = boost::get<musrfit::ast::expression>(&expr.first)) {

    // Check if inner expression has low-precedence ops that need parentheses
    bool needsParens = false;
    if (hasHighPrecOps && !inner->rest.empty()) {
      for (const auto& innerOp : inner->rest) {
        if (innerOp.operator_ == musrfit::ast::op_plus || innerOp.operator_ == musrfit::ast::op_minus) {
          needsParens = true;
          break;
        }
      }
    }

    if (needsParens) {
      result = "(" + GenerateString(*inner) + ")";
    } else {
      result = GenerateString(*inner);
    }
  } else {
    result = GenerateStringOperand(expr.first);
  }

  // Generate operations and operands
  for (const auto& op : expr.rest) {
    switch (op.operator_) {
      case musrfit::ast::op_plus:  result += " + "; break;
      case musrfit::ast::op_minus: result += " - "; break;
      case musrfit::ast::op_times: result += " * "; break;
      case musrfit::ast::op_divide: result += " / "; break;
    }

    // Check if operand is a wrapped expression
    if (const musrfit::ast::expression* inner = boost::get<musrfit::ast::expression>(&op.operand_)) {

      // Check if inner expression has low-precedence ops that need parentheses
      bool needsParens = false;
      if ((op.operator_ == musrfit::ast::op_times || op.operator_ == musrfit::ast::op_divide) &&
          !inner->rest.empty()) {
        for (const auto& innerOp : inner->rest) {
          if (innerOp.operator_ == musrfit::ast::op_plus || innerOp.operator_ == musrfit::ast::op_minus) {
            needsParens = true;
            break;
          }
        }
      }

      if (needsParens) {
        result += "(" + GenerateString(*inner) + ")";
      } else {
        result += GenerateString(*inner);
      }
    } else {
      result += GenerateStringOperand(op.operand_);
    }
  }

  return result;
}

//--------------------------------------------------------------------------
/**
 * \brief Generates a string representation of an operand.
 *
 * \param operand Operand AST to convert
 * \return Formatted string
 */
TString PFunction::GenerateStringOperand(const musrfit::ast::operand& operand)
{
  struct StringVisitor : public boost::static_visitor<TString>
  {
    PFunction* fFunc;
    StringVisitor(PFunction* func) : fFunc(func) {}

    TString operator()(const musrfit::ast::nil&) const { return "nil"; }

    TString operator()(double val) const {
      std::ostringstream oss;
      oss << val;
      return TString(oss.str().c_str());
    }

    TString operator()(const musrfit::ast::constant& c) const {
      TString str;
      if (c.sign) str = "-";

      switch (c.const_type) {
        case musrfit::ast::constant::pi:       str += "PI"; break;
        case musrfit::ast::constant::gamma_mu: str += "GAMMA_MU"; break;
        case musrfit::ast::constant::field:    str += (c.sign ? "B" : "B"); break;  // Sign already added
        case musrfit::ast::constant::energy:   str += (c.sign ? "EN" : "EN"); break;
        case musrfit::ast::constant::temp:
          str += "T";
          str += c.index;
          break;
      }
      return str;
    }

    TString operator()(const musrfit::ast::parameter& p) const {
      TString str;
      if (p.sign) str = "-";
      str += "PAR";
      str += p.number;
      return str;
    }

    TString operator()(const musrfit::ast::map_ref& m) const {
      TString str;
      if (m.sign) str = "-";
      str += "MAP";
      str += m.number;
      return str;
    }

    TString operator()(const musrfit::ast::function_call& f) const {
      TString str;
      switch (f.func_id) {
        case musrfit::ast::fun_cos:   str = "COS"; break;
        case musrfit::ast::fun_sin:   str = "SIN"; break;
        case musrfit::ast::fun_tan:   str = "TAN"; break;
        case musrfit::ast::fun_cosh:  str = "COSH"; break;
        case musrfit::ast::fun_sinh:  str = "SINH"; break;
        case musrfit::ast::fun_tanh:  str = "TANH"; break;
        case musrfit::ast::fun_acos:  str = "ACOS"; break;
        case musrfit::ast::fun_asin:  str = "ASIN"; break;
        case musrfit::ast::fun_atan:  str = "ATAN"; break;
        case musrfit::ast::fun_acosh: str = "ACOSH"; break;
        case musrfit::ast::fun_asinh: str = "ASINH"; break;
        case musrfit::ast::fun_atanh: str = "ATANH"; break;
        case musrfit::ast::fun_log:   str = "LOG"; break;
        case musrfit::ast::fun_ln:    str = "LN"; break;
        case musrfit::ast::fun_exp:   str = "EXP"; break;
        case musrfit::ast::fun_sqrt:  str = "SQRT"; break;
      }
      str += "(";
      str += fFunc->GenerateString(f.arg);
      str += ")";
      return str;
    }

    TString operator()(const musrfit::ast::power_call& p) const {
      TString str = "POW(";
      str += fFunc->GenerateString(p.base);
      str += ", ";
      str += fFunc->GenerateString(p.pow);
      str += ")";
      return str;
    }

    TString operator()(const musrfit::ast::expression& expr) const {
      // If no operations, just return the first operand without parentheses
      if (expr.rest.empty()) {
        return fFunc->GenerateStringOperand(expr.first);
      } else {
        // Has operations - this means we have a parenthesized expression from the grammar
        // (e.g., from factor = '(' expression ')'), so we need to preserve the parentheses
        return "(" + fFunc->GenerateString(expr) + ")";
      }
    }
  };

  StringVisitor visitor(this);
  return boost::apply_visitor(visitor, operand);
}

//==========================================================================
// EvalVisitor Implementation
//==========================================================================

/**
 * \brief Evaluates a nil (empty) AST node.
 *
 * \return Always returns 0.0
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::nil&) const
{
  return 0.0;
}

/**
 * \brief Evaluates a numeric literal.
 *
 * \param val The literal value from the AST
 * \return The literal value unchanged
 */
Double_t PFunction::EvalVisitor::operator()(double val) const
{
  return val;
}

/**
 * \brief Evaluates a constant (PI, GAMMA_MU, B, EN, T#).
 *
 * Converts symbolic constants to their numeric values, including full support
 * for experimental metadata (magnetic field, energy, temperature).
 *
 * \param c The constant AST node with type, sign, and optional index
 * \return The constant's numeric value (with sign applied if negative)
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::constant& c) const
{
  Double_t val = 0.0;

  switch (c.const_type) {
    case musrfit::ast::constant::pi:
      val = 3.14159265358979323846;
      break;

    case musrfit::ast::constant::gamma_mu:
      val = GAMMA_BAR_MUON;  // From PMusr.h
      break;

    case musrfit::ast::constant::field:
      val = fMetaData.fField;  // Magnetic field (Gauss)
      break;

    case musrfit::ast::constant::energy:
      val = fMetaData.fEnergy;  // Energy (keV)
      break;

    case musrfit::ast::constant::temp:
      // Temperature array (Kelvin)
      if (c.index >= 0 && static_cast<UInt_t>(c.index) < fMetaData.fTemp.size()) {
        val = fMetaData.fTemp[c.index];
      } else {
        std::cerr << "**ERROR** Temperature index T" << c.index << " out of range" << std::endl;
        val = 0.0;
      }
      break;

    default:
      val = 0.0;
  }

  return c.sign ? -val : val;
}

/**
 * \brief Evaluates a parameter reference (PAR#).
 *
 * Looks up the parameter value from the fParam vector using the
 * parameter number (1-based indexing in syntax, 0-based in vector).
 * Validates bounds and applies sign.
 *
 * \param p The parameter AST node containing number and sign
 * \return The parameter value (with sign applied), or 0.0 if out of bounds
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::parameter& p) const
{
  Int_t idx = p.number - 1;

  if (idx < 0 || static_cast<UInt_t>(idx) >= fParam.size()) {
    std::cerr << "**ERROR** Parameter PAR" << p.number << " out of range during evaluation" << std::endl;
    return 0.0;
  }

  Double_t val = fParam[idx];
  return p.sign ? -val : val;
}

/**
 * \brief Evaluates a map reference (MAP#) for indirect parameter lookup.
 *
 * Performs a two-level lookup: first looks up the map index in fMap,
 * then uses that value to index into fParam. This allows for flexible
 * parameter remapping. Validates both lookups and applies sign.
 *
 * \param m The map reference AST node containing map number and sign
 * \return The indirectly referenced parameter value (with sign), or 0.0 if invalid
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::map_ref& m) const
{
  Int_t mapIdx = m.number - 1;

  if (mapIdx < 0 || static_cast<UInt_t>(mapIdx) >= fMap.size()) {
    std::cerr << "**ERROR** Map MAP" << m.number << " out of range during evaluation" << std::endl;
    return 0.0;
  }

  Int_t paramIdx = fMap[mapIdx] - 1;

  if (paramIdx < 0 || static_cast<UInt_t>(paramIdx) >= fParam.size()) {
    std::cerr << "**ERROR** Mapped parameter out of range during evaluation" << std::endl;
    return 0.0;
  }

  Double_t val = fParam[paramIdx];
  return m.sign ? -val : val;
}

/**
 * \brief Evaluates a function call (COS, SIN, EXP, SQRT, etc.).
 *
 * Recursively evaluates the argument expression, then applies the
 * specified mathematical function. Supports trigonometric, hyperbolic,
 * inverse, exponential, and logarithmic functions.
 *
 * \param f The function call AST node with function ID and argument expression
 * \return The result of applying the mathematical function to the argument
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::function_call& f) const
{
  // Evaluate argument
  Double_t arg = boost::apply_visitor(*this, f.arg.first);
  for (const auto& op : f.arg.rest) {
    Double_t rhs = boost::apply_visitor(*this, op.operand_);
    switch (op.operator_) {
      case musrfit::ast::op_plus:   arg += rhs; break;
      case musrfit::ast::op_minus:  arg -= rhs; break;
      case musrfit::ast::op_times:  arg *= rhs; break;
      case musrfit::ast::op_divide:
        if (std::abs(rhs) > 1.0e-20) arg /= rhs;
        break;
    }
  }

  // Apply function
  Double_t result = 0.0;
  switch (f.func_id) {
    case musrfit::ast::fun_cos:   result = std::cos(arg); break;
    case musrfit::ast::fun_sin:   result = std::sin(arg); break;
    case musrfit::ast::fun_tan:   result = std::tan(arg); break;
    case musrfit::ast::fun_cosh:  result = std::cosh(arg); break;
    case musrfit::ast::fun_sinh:  result = std::sinh(arg); break;
    case musrfit::ast::fun_tanh:  result = std::tanh(arg); break;
    case musrfit::ast::fun_acos:  result = std::acos(arg); break;
    case musrfit::ast::fun_asin:  result = std::asin(arg); break;
    case musrfit::ast::fun_atan:  result = std::atan(arg); break;
    case musrfit::ast::fun_acosh: result = std::acosh(arg); break;
    case musrfit::ast::fun_asinh: result = std::asinh(arg); break;
    case musrfit::ast::fun_atanh: result = std::atanh(arg); break;
    case musrfit::ast::fun_log:   result = std::log10(std::abs(arg)); break;
    case musrfit::ast::fun_ln:    result = std::log(std::abs(arg)); break;
    case musrfit::ast::fun_exp:   result = std::exp(arg); break;
    case musrfit::ast::fun_sqrt:  result = std::sqrt(std::abs(arg)); break;
    default: result = 0.0;
  }

  return result;
}

/**
 * \brief Evaluates a power operation POW(base, exponent).
 *
 * Recursively evaluates both the base and exponent expressions,
 * then computes base^exponent using std::pow. Uses absolute value
 * of base to avoid complex number issues with negative bases.
 *
 * \param p The power call AST node with base and exponent expressions
 * \return The result of raising the base to the exponent power
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::power_call& p) const
{
  // Evaluate base
  Double_t base = boost::apply_visitor(*this, p.base.first);
  for (const auto& op : p.base.rest) {
    Double_t rhs = boost::apply_visitor(*this, op.operand_);
    switch (op.operator_) {
      case musrfit::ast::op_plus:   base += rhs; break;
      case musrfit::ast::op_minus:  base -= rhs; break;
      case musrfit::ast::op_times:  base *= rhs; break;
      case musrfit::ast::op_divide:
        if (std::abs(rhs) > 1.0e-20) base /= rhs;
        break;
    }
  }

  // Evaluate exponent
  Double_t exponent = boost::apply_visitor(*this, p.pow.first);
  for (const auto& op : p.pow.rest) {
    Double_t rhs = boost::apply_visitor(*this, op.operand_);
    switch (op.operator_) {
      case musrfit::ast::op_plus:   exponent += rhs; break;
      case musrfit::ast::op_minus:  exponent -= rhs; break;
      case musrfit::ast::op_times:  exponent *= rhs; break;
      case musrfit::ast::op_divide:
        if (std::abs(rhs) > 1.0e-20) exponent /= rhs;
        break;
    }
  }

  return std::pow(std::abs(base), exponent);
}

/**
 * \brief Evaluates a complete expression with binary operators.
 *
 * Evaluates the first operand, then applies each operation in sequence
 * from left to right. The AST structure encodes operator precedence,
 * so this left-to-right evaluation produces correct results.
 *
 * Handles division by zero by checking if divisor magnitude is greater
 * than 1e-20; otherwise returns 0.0 for that sub-expression.
 *
 * \param e The expression AST node with first operand and operation list
 * \return The computed result of the expression
 */
Double_t PFunction::EvalVisitor::operator()(const musrfit::ast::expression& e) const
{
  Double_t result = boost::apply_visitor(*this, e.first);

  for (const auto& op : e.rest) {
    Double_t rhs = boost::apply_visitor(*this, op.operand_);
    switch (op.operator_) {
      case musrfit::ast::op_plus:   result += rhs; break;
      case musrfit::ast::op_minus:  result -= rhs; break;
      case musrfit::ast::op_times:  result *= rhs; break;
      case musrfit::ast::op_divide:
        if (std::abs(rhs) > 1.0e-20) result /= rhs;
        break;
    }
  }

  return result;
}