parser.go

package mathopt

import (
	"fmt"
	"strconv"
	"unicode"
)

// TokenType represents the type of a token
type TokenType int

const (
	TokenNumber TokenType = iota
	TokenVariable
	TokenOperator
	TokenLeftParen
	TokenRightParen
	TokenFunction
	TokenEOF
)

// Token represents a lexical token
type Token struct {
	Type  TokenType
	Value string
}

// Tokenizer converts a string into tokens
type Tokenizer struct {
	input string
	pos   int
}

// NewTokenizer creates a new tokenizer
func NewTokenizer(input string) *Tokenizer {
	return &Tokenizer{input: input, pos: 0}
}

// NextToken returns the next token from the input
func (t *Tokenizer) NextToken() (*Token, error) {
	// Skip whitespace
	for t.pos < len(t.input) && unicode.IsSpace(rune(t.input[t.pos])) {
		t.pos++
	}
	
	if t.pos >= len(t.input) {
		return &Token{Type: TokenEOF}, nil
	}
	
	ch := t.input[t.pos]
	
	// Number
	if unicode.IsDigit(rune(ch)) || ch == '.' {
		return t.readNumber()
	}
	
	// Variable or function
	if unicode.IsLetter(rune(ch)) {
		return t.readIdentifier()
	}
	
	// Operators and parentheses
	t.pos++
	switch ch {
	case '+', '*', '/', '^':
		return &Token{Type: TokenOperator, Value: string(ch)}, nil
	case '-':
		// Could be unary minus or binary minus
		return &Token{Type: TokenOperator, Value: string(ch)}, nil
	case '(':
		return &Token{Type: TokenLeftParen, Value: string(ch)}, nil
	case ')':
		return &Token{Type: TokenRightParen, Value: string(ch)}, nil
	default:
		return nil, fmt.Errorf("unexpected character: %c at position %d", ch, t.pos-1)
	}
}

func (t *Tokenizer) readNumber() (*Token, error) {
	start := t.pos
	hasDot := false
	
	for t.pos < len(t.input) {
		ch := t.input[t.pos]
		if unicode.IsDigit(rune(ch)) {
			t.pos++
		} else if ch == '.' && !hasDot {
			hasDot = true
			t.pos++
		} else {
			break
		}
	}
	
	value := t.input[start:t.pos]
	return &Token{Type: TokenNumber, Value: value}, nil
}

func (t *Tokenizer) readIdentifier() (*Token, error) {
	start := t.pos
	
	for t.pos < len(t.input) && (unicode.IsLetter(rune(t.input[t.pos])) || unicode.IsDigit(rune(t.input[t.pos]))) {
		t.pos++
	}
	
	value := t.input[start:t.pos]
	
	// Check if it's a function
	functions := []string{"sin", "cos", "exp", "ln"}
	for _, fn := range functions {
		if value == fn {
			return &Token{Type: TokenFunction, Value: value}, nil
		}
	}
	
	return &Token{Type: TokenVariable, Value: value}, nil
}

// Parser parses tokens into an expression tree
type Parser struct {
	tokens  []*Token
	current int
}

// Parse parses a mathematical expression string into an Expr
func Parse(input string) (Expr, error) {
	tokenizer := NewTokenizer(input)
	var tokens []*Token
	
	for {
		token, err := tokenizer.NextToken()
		if err != nil {
			return nil, err
		}
		tokens = append(tokens, token)
		if token.Type == TokenEOF {
			break
		}
	}
	
	parser := &Parser{tokens: tokens, current: 0}
	return parser.parseExpression()
}

func (p *Parser) parseExpression() (Expr, error) {
	return p.parseAddSub()
}

func (p *Parser) parseAddSub() (Expr, error) {
	left, err := p.parseMulDiv()
	if err != nil {
		return nil, err
	}
	
	for p.current < len(p.tokens) {
		token := p.tokens[p.current]
		if token.Type == TokenOperator && (token.Value == "+" || token.Value == "-") {
			p.current++
			right, err := p.parseMulDiv()
			if err != nil {
				return nil, err
			}
			left = &BinaryOp{Op: token.Value, Left: left, Right: right}
		} else {
			break
		}
	}
	
	return left, nil
}

func (p *Parser) parseMulDiv() (Expr, error) {
	left, err := p.parsePower()
	if err != nil {
		return nil, err
	}
	
	for p.current < len(p.tokens) {
		token := p.tokens[p.current]
		if token.Type == TokenOperator && (token.Value == "*" || token.Value == "/") {
			p.current++
			right, err := p.parsePower()
			if err != nil {
				return nil, err
			}
			left = &BinaryOp{Op: token.Value, Left: left, Right: right}
		} else {
			break
		}
	}
	
	return left, nil
}

func (p *Parser) parsePower() (Expr, error) {
	left, err := p.parseUnary()
	if err != nil {
		return nil, err
	}
	
	for p.current < len(p.tokens) {
		token := p.tokens[p.current]
		if token.Type == TokenOperator && token.Value == "^" {
			p.current++
			right, err := p.parseUnary()
			if err != nil {
				return nil, err
			}
			left = &BinaryOp{Op: token.Value, Left: left, Right: right}
		} else {
			break
		}
	}
	
	return left, nil
}

func (p *Parser) parseUnary() (Expr, error) {
	if p.current >= len(p.tokens) {
		return nil, fmt.Errorf("unexpected end of expression")
	}
	
	token := p.tokens[p.current]
	
	// Unary minus
	if token.Type == TokenOperator && token.Value == "-" {
		p.current++
		arg, err := p.parseUnary()
		if err != nil {
			return nil, err
		}
		return &UnaryOp{Op: "-", Arg: arg}, nil
	}
	
	// Function call
	if token.Type == TokenFunction {
		fnName := token.Value
		p.current++
		
		if p.current >= len(p.tokens) || p.tokens[p.current].Type != TokenLeftParen {
			return nil, fmt.Errorf("expected '(' after function %s", fnName)
		}
		p.current++ // skip '('
		
		arg, err := p.parseExpression()
		if err != nil {
			return nil, err
		}
		
		if p.current >= len(p.tokens) || p.tokens[p.current].Type != TokenRightParen {
			return nil, fmt.Errorf("expected ')' after function argument")
		}
		p.current++ // skip ')'
		
		return &UnaryOp{Op: fnName, Arg: arg}, nil
	}
	
	return p.parsePrimary()
}

func (p *Parser) parsePrimary() (Expr, error) {
	if p.current >= len(p.tokens) {
		return nil, fmt.Errorf("unexpected end of expression")
	}
	
	token := p.tokens[p.current]
	p.current++
	
	switch token.Type {
	case TokenNumber:
		value, err := strconv.ParseFloat(token.Value, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid number: %s", token.Value)
		}
		return &Constant{Value: value}, nil
		
	case TokenVariable:
		return &Variable{Name: token.Value}, nil
		
	case TokenLeftParen:
		expr, err := p.parseExpression()
		if err != nil {
			return nil, err
		}
		
		if p.current >= len(p.tokens) || p.tokens[p.current].Type != TokenRightParen {
			return nil, fmt.Errorf("expected ')'")
		}
		p.current++
		return expr, nil
		
	default:
		return nil, fmt.Errorf("unexpected token: %s", token.Value)
	}
}

// MustParse parses an expression and panics on error (useful for tests)
func MustParse(input string) Expr {
	expr, err := Parse(input)
	if err != nil {
		panic(err)
	}
	return expr
}

// ExpressionEquivalence checks if two expressions are equivalent
// by simplifying both and comparing them structurally
func ExpressionEquivalence(e1, e2 Expr) bool {
	s1 := e1.Simplify()
	s2 := e2.Simplify()
	return s1.Equals(s2)
}

// NormalizeExpression applies algebraic normalization to an expression
// This includes sorting commutative operations, collecting like terms, etc.
func NormalizeExpression(expr Expr) Expr {
	simplified := expr.Simplify()
	// Additional normalization could be added here
	return simplified
}

// ConstantFold performs aggressive constant folding on an expression
func ConstantFold(expr Expr) Expr {
	// Simplify already does constant folding, but we can be more aggressive
	prev := expr
	for {
		current := prev.Simplify()
		if current.Equals(prev) {
			break
		}
		prev = current
	}
	return prev
}

// PrettyPrint returns a nicely formatted string representation
func PrettyPrint(expr Expr) string {
	return expr.String()
}