#include "expr.h"

#include <stdlib.h>
#include <string.h>

static Expr *expr_node(ExprKind kind) {
    Expr *e = malloc(sizeof *e);
    if (e == NULL) {
        return NULL;
    }
    e->kind = kind;
    return e;
}

Expr *expr_number(double value) {
    Expr *e = expr_node(EXPR_NUMBER);
    if (e == NULL) {
        return NULL;
    }
    e->as.number = value;
    return e;
}

Expr *expr_variable(const char *name) {
    Expr *e = expr_node(EXPR_VARIABLE);
    if (e == NULL) {
        return NULL;
    }
    e->as.name = malloc(strlen(name) + 1);
    if (e->as.name == NULL) {
        free(e);
        return NULL;
    }
    strcpy(e->as.name, name);
    return e;
}

Expr *expr_binary(BinaryOp op, Expr *left, Expr *right) {
    if (left == NULL || right == NULL) {
        expr_free(left);
        expr_free(right);
        return NULL;
    }
    Expr *e = expr_node(EXPR_BINARY);
    if (e == NULL) {
        expr_free(left);
        expr_free(right);
        return NULL;
    }
    e->as.binary.op = op;
    e->as.binary.left = left;
    e->as.binary.right = right;
    return e;
}

void expr_free(Expr *e) {
    if (e == NULL) {
        return;
    }
    switch (e->kind) {
    case EXPR_NUMBER:
        break;
    case EXPR_VARIABLE:
        free(e->as.name);
        break;
    case EXPR_BINARY:
        expr_free(e->as.binary.left);
        expr_free(e->as.binary.right);
        break;
    }
    free(e);
}

// El nivel de las comparaciones, el más bajo.
enum { COMPARISON = 1 };

// Mayor número, liga más fuerte. Las hojas nunca necesitan paréntesis.
static int precedence(const Expr *e) {
    if (e->kind != EXPR_BINARY) {
        return 4;
    }
    switch (e->as.binary.op) {
    case OP_MUL:
    case OP_DIV:
        return 3;
    case OP_ADD:
    case OP_SUB:
        return 2;
    case OP_LESS:
    case OP_LESS_EQUAL:
    case OP_GREATER:
    case OP_GREATER_EQUAL:
    case OP_EQUAL:
    case OP_NOT_EQUAL:
        break;
    }
    return COMPARISON;
}

static const char *op_symbol(BinaryOp op) {
    switch (op) {
    case OP_ADD:
        return "+";
    case OP_SUB:
        return "-";
    case OP_MUL:
        return "*";
    case OP_DIV:
        return "/";
    case OP_LESS:
        return "<";
    case OP_LESS_EQUAL:
        return "<=";
    case OP_GREATER:
        return ">";
    case OP_GREATER_EQUAL:
        return ">=";
    case OP_EQUAL:
        return "==";
    case OP_NOT_EQUAL:
        return "!=";
    }
    return "?";
}

static void print_child(const Expr *child, int needs_parens, FILE *out) {
    if (needs_parens) {
        fputc('(', out);
    }
    expr_print(child, out);
    if (needs_parens) {
        fputc(')', out);
    }
}

void expr_print(const Expr *e, FILE *out) {
    switch (e->kind) {
    case EXPR_NUMBER:
        fprintf(out, "%g", e->as.number);
        break;
    case EXPR_VARIABLE:
        fputs(e->as.name, out);
        break;
    case EXPR_BINARY: {
        const Expr *left = e->as.binary.left;
        const Expr *right = e->as.binary.right;
        // Las comparaciones no se encadenan: a la izquierda, el mismo nivel también lleva paréntesis.
        int level = precedence(e);
        print_child(left, precedence(left) < level || (level == COMPARISON && precedence(left) == level), out);
        fprintf(out, " %s ", op_symbol(e->as.binary.op));
        print_child(right, precedence(right) <= level, out);
        break;
    }
    }
}

// Aplica op a dos números. Devuelve 0 si es una división entre 0.
static int apply_op(BinaryOp op, double left, double right, double *result) {
    switch (op) {
    case OP_ADD:
        *result = left + right;
        return 1;
    case OP_SUB:
        *result = left - right;
        return 1;
    case OP_MUL:
        *result = left * right;
        return 1;
    case OP_DIV:
        if (right == 0.0) {
            return 0;
        }
        *result = left / right;
        return 1;
    case OP_LESS:
        *result = left < right;
        return 1;
    case OP_LESS_EQUAL:
        *result = left <= right;
        return 1;
    case OP_GREATER:
        *result = left > right;
        return 1;
    case OP_GREATER_EQUAL:
        *result = left >= right;
        return 1;
    case OP_EQUAL:
        *result = left == right;
        return 1;
    case OP_NOT_EQUAL:
        *result = left != right;
        return 1;
    }
    return 0;
}

int expr_eval(const Expr *e, const Env *env, double *result) {
    switch (e->kind) {
    case EXPR_NUMBER:
        *result = e->as.number;
        return 1;
    case EXPR_VARIABLE:
        return env_get(env, e->as.name, result);
    case EXPR_BINARY: {
        double left;
        double right;
        if (!expr_eval(e->as.binary.left, env, &left)) {
            return 0;
        }
        if (!expr_eval(e->as.binary.right, env, &right)) {
            return 0;
        }
        return apply_op(e->as.binary.op, left, right, result);
    }
    }
    return 0;
}

Expr *expr_copy(const Expr *e) {
    switch (e->kind) {
    case EXPR_NUMBER:
        return expr_number(e->as.number);
    case EXPR_VARIABLE:
        return expr_variable(e->as.name);
    case EXPR_BINARY:
        return expr_binary(e->as.binary.op, expr_copy(e->as.binary.left), expr_copy(e->as.binary.right));
    }
    return NULL;
}

int expr_equal(const Expr *a, const Expr *b) {
    if (a->kind != b->kind) {
        return 0;
    }
    switch (a->kind) {
    case EXPR_NUMBER:
        return a->as.number == b->as.number;
    case EXPR_VARIABLE:
        return strcmp(a->as.name, b->as.name) == 0;
    case EXPR_BINARY:
        return a->as.binary.op == b->as.binary.op
            && expr_equal(a->as.binary.left, b->as.binary.left)
            && expr_equal(a->as.binary.right, b->as.binary.right);
    }
    return 0;
}

static int is_number(const Expr *e, double value) {
    return e->kind == EXPR_NUMBER && e->as.number == value;
}

// Conserva un hijo, y libera al otro hijo y al nodo.
static Expr *keep_child(Expr *node, Expr *kept, Expr *dropped) {
    expr_free(dropped);
    free(node);
    return kept;
}

// Convierte un nodo binario en un número, en el mismo lugar. No reserva nada.
static Expr *become_number(Expr *node, double value) {
    expr_free(node->as.binary.left);
    expr_free(node->as.binary.right);
    node->kind = EXPR_NUMBER;
    node->as.number = value;
    return node;
}

Expr *expr_simplify(Expr *e) {
    if (e == NULL || e->kind != EXPR_BINARY) {
        return e;
    }
    e->as.binary.left = expr_simplify(e->as.binary.left);
    e->as.binary.right = expr_simplify(e->as.binary.right);
    Expr *left = e->as.binary.left;
    Expr *right = e->as.binary.right;
    BinaryOp op = e->as.binary.op;

    if ((op == OP_ADD && is_number(left, 0.0)) || (op == OP_MUL && is_number(left, 1.0))) {
        return keep_child(e, right, left);
    }
    if ((op == OP_ADD || op == OP_SUB) && is_number(right, 0.0)) {
        return keep_child(e, left, right);
    }
    if ((op == OP_MUL || op == OP_DIV) && is_number(right, 1.0)) {
        return keep_child(e, left, right);
    }
    if (op == OP_MUL && (is_number(left, 0.0) || is_number(right, 0.0))) {
        return become_number(e, 0.0);
    }
    if (left->kind == EXPR_NUMBER && right->kind == EXPR_NUMBER) {
        double value;
        if (apply_op(op, left->as.number, right->as.number, &value)) {
            return become_number(e, value);
        }
    }
    return e;
}

Expr *expr_derive(const Expr *e, const char *name) {
    switch (e->kind) {
    case EXPR_NUMBER:
        return expr_number(0.0);
    case EXPR_VARIABLE:
        return expr_number(strcmp(e->as.name, name) == 0 ? 1.0 : 0.0);
    case EXPR_BINARY: {
        const Expr *f = e->as.binary.left;
        const Expr *g = e->as.binary.right;
        switch (e->as.binary.op) {
        case OP_ADD:
        case OP_SUB:
            return expr_binary(e->as.binary.op, expr_derive(f, name), expr_derive(g, name));
        case OP_MUL:
            return expr_binary(OP_ADD,
                               expr_binary(OP_MUL, expr_derive(f, name), expr_copy(g)),
                               expr_binary(OP_MUL, expr_copy(f), expr_derive(g, name)));
        case OP_DIV:
            return expr_binary(OP_DIV,
                               expr_binary(OP_SUB,
                                           expr_binary(OP_MUL, expr_derive(f, name), expr_copy(g)),
                                           expr_binary(OP_MUL, expr_copy(f), expr_derive(g, name))),
                               expr_binary(OP_MUL, expr_copy(g), expr_copy(g)));
        case OP_LESS:
        case OP_LESS_EQUAL:
        case OP_GREATER:
        case OP_GREATER_EQUAL:
        case OP_EQUAL:
        case OP_NOT_EQUAL:
            return NULL;
        }
        return NULL;
    }
    }
    return NULL;
}