// probol_repl.c – PROBOL Interactive REPL with Standard Library
// Compile: gcc -o probol probol_repl.c -lm
// Run: ./probol

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <math.h>
#include <setjmp.h>
#include <signal.h>

// ---------- Forward declarations ----------
static void trim(char* s);
static void process_line(char* line);

// ---------- Probability Tensor ----------
#define MAX_STATES 64
typedef struct {
    double value;
    double prob;
} State;

typedef struct {
    State states[MAX_STATES];
    int count;
} Tensor;

// ---------- Symbol Table ----------
typedef enum { SYM_TENSOR, SYM_SCALAR } SymType;
typedef struct {
    char name[64];
    SymType type;
    union {
        Tensor t;
        double d;
    } value;
} Symbol;

static Symbol symtab[256];
static int sym_count = 0;

static void add_tensor(const char* name, Tensor* t) {
    strcpy(symtab[sym_count].name, name);
    symtab[sym_count].type = SYM_TENSOR;
    symtab[sym_count].value.t = *t;
    sym_count++;
}

static void add_scalar(const char* name, double val) {
    strcpy(symtab[sym_count].name, name);
    symtab[sym_count].type = SYM_SCALAR;
    symtab[sym_count].value.d = val;
    sym_count++;
}

static Symbol* find_sym(const char* name) {
    for (int i = 0; i < sym_count; i++)
        if (strcmp(symtab[i].name, name) == 0)
            return &symtab[i];
    return NULL;
}

// ---------- Tensor Operations ----------
static void init_tensor(Tensor* t) { t->count = 0; }

static void add_state(Tensor* t, double val, double prob) {
    if (t->count >= MAX_STATES) return;
    // merge equal values
    for (int i = 0; i < t->count; i++) {
        if (fabs(t->states[i].value - val) < 1e-12) {
            t->states[i].prob += prob;
            return;
        }
    }
    t->states[t->count].value = val;
    t->states[t->count].prob = prob;
    t->count++;
}

static void normalize(Tensor* t) {
    double sum = 0;
    for (int i = 0; i < t->count; i++) sum += t->states[i].prob;
    if (sum > 0)
        for (int i = 0; i < t->count; i++) t->states[i].prob /= sum;
}

static void conv_add(Tensor* a, Tensor* b, Tensor* res) {
    init_tensor(res);
    for (int i = 0; i < a->count; i++)
        for (int j = 0; j < b->count; j++) {
            double val = a->states[i].value + b->states[j].value;
            double prob = a->states[i].prob * b->states[j].prob;
            add_state(res, val, prob);
        }
    normalize(res);
}

static double measure(Tensor* t) {
    if (t->count == 0) return 0;
    int best = 0;
    double maxp = t->states[0].prob;
    for (int i = 1; i < t->count; i++)
        if (t->states[i].prob > maxp) { maxp = t->states[i].prob; best = i; }
    return t->states[best].value;
}

static double mean(Tensor* t) {
    double m = 0;
    for (int i = 0; i < t->count; i++) m += t->states[i].value * t->states[i].prob;
    return m;
}

static double variance(Tensor* t) {
    double m = mean(t);
    double var = 0;
    for (int i = 0; i < t->count; i++) {
        double diff = t->states[i].value - m;
        var += diff * diff * t->states[i].prob;
    }
    return var;
}

static double quantile(Tensor* t, double p) {
    if (t->count == 0) return 0;
    // sort by value
    State sorted[MAX_STATES];
    memcpy(sorted, t->states, sizeof(State)*t->count);
    for (int i = 0; i < t->count-1; i++)
        for (int j = i+1; j < t->count; j++)
            if (sorted[i].value > sorted[j].value) {
                State tmp = sorted[i]; sorted[i] = sorted[j]; sorted[j] = tmp;
            }
    double cum = 0;
    for (int i = 0; i < t->count; i++) {
        cum += sorted[i].prob;
        if (cum >= p) return sorted[i].value;
    }
    return sorted[t->count-1].value;
}

static void print_tensor(Tensor* t) {
    printf("Tensor (%d states):\n", t->count);
    for (int i = 0; i < t->count; i++)
        printf("  %.6f : %.6f\n", t->states[i].value, t->states[i].prob);
}

// ---------- Standard Library Distributions ----------
static Tensor gaussian(double mu, double sigma) {
    Tensor t; init_tensor(&t);
    // approximate with 5 sigma points
    double vals[] = {mu-2*sigma, mu-sigma, mu, mu+sigma, mu+2*sigma};
    double probs[] = {0.054, 0.242, 0.408, 0.242, 0.054};
    for (int i=0; i<5; i++) add_state(&t, vals[i], probs[i]);
    normalize(&t);
    return t;
}

static Tensor uniform(double a, double b) {
    Tensor t; init_tensor(&t);
    // simple discrete approximation
    double mid = (a+b)/2;
    add_state(&t, mid, 1.0);
    return t;
}

static Tensor binomial(int n, double p) {
    Tensor t; init_tensor(&t);
    double mean = n * p;
    double var = n * p * (1-p);
    double sd = sqrt(var);
    int v0 = (int)round(mean);
    if (v0 < 0) v0 = 0; if (v0 > n) v0 = n;
    add_state(&t, v0, 0.5);
    if (sd > 0.1) {
        int v1 = v0 + 1; if (v1 <= n) add_state(&t, v1, 0.25);
        int v2 = v0 - 1; if (v2 >= 0) add_state(&t, v2, 0.25);
    }
    normalize(&t);
    return t;
}

static Tensor exponential(double lambda) {
    Tensor t; init_tensor(&t);
    double mean = 1.0/lambda;
    add_state(&t, mean, 1.0);
    return t;
}

static Tensor poisson(double lambda) {
    Tensor t; init_tensor(&t);
    int k = (int)round(lambda);
    add_state(&t, k, 1.0);
    return t;
}

static Tensor lognormal(double mu, double sigma) {
    Tensor t; init_tensor(&t);
    double m = exp(mu + sigma*sigma/2);
    add_state(&t, m, 1.0);
    return t;
}

// ---------- Expression Evaluation (simple recursive descent) ----------
static char* expr_str;
static int expr_pos;
static double expr_eval();
static double expr_primary();
static double expr_factor();
static double expr_term();

static double expr_eval() { return expr_term(); }
static double expr_term() {
    double left = expr_factor();
    while (1) {
        while (isspace(expr_str[expr_pos])) expr_pos++;
        char op = expr_str[expr_pos];
        if (op == '+' || op == '-') {
            expr_pos++;
            double right = expr_factor();
            if (op == '+') left += right;
            else left -= right;
        } else break;
    }
    return left;
}
static double expr_factor() {
    double left = expr_primary();
    while (1) {
        while (isspace(expr_str[expr_pos])) expr_pos++;
        char op = expr_str[expr_pos];
        if (op == '*' || op == '/') {
            expr_pos++;
            double right = expr_primary();
            if (op == '*') left *= right;
            else left /= right;
        } else break;
    }
    return left;
}
static double expr_primary() {
    while (isspace(expr_str[expr_pos])) expr_pos++;
    if (expr_str[expr_pos] == '(') {
        expr_pos++;
        double val = expr_eval();
        while (isspace(expr_str[expr_pos])) expr_pos++;
        if (expr_str[expr_pos] == ')') expr_pos++;
        return val;
    }
    if (isalpha(expr_str[expr_pos])) {
        char name[64]; int i=0;
        while (isalnum(expr_str[expr_pos]) || expr_str[expr_pos]=='_')
            name[i++] = expr_str[expr_pos++];
        name[i]=0;
        if (strcmp(name, "sqrt") == 0) {
            while (isspace(expr_str[expr_pos])) expr_pos++;
            if (expr_str[expr_pos] == '(') expr_pos++;
            double arg = expr_eval();
            while (isspace(expr_str[expr_pos])) expr_pos++;
            if (expr_str[expr_pos] == ')') expr_pos++;
            return sqrt(arg);
        }
        // variable lookup
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_SCALAR) return s->value.d;
        else {
            printf("Undefined scalar: %s\n", name);
            return 0;
        }
    }
    double val;
    int n = 0;
    sscanf(expr_str+expr_pos, "%lf%n", &val, &n);
    expr_pos += n;
    return val;
}

// ---------- PROBOL Line Parsing and Execution ----------
static void process_line(char* line) {
    trim(line);
    if (strlen(line)==0 || line[0]==';') return;
    
    // dist name = { val:prob, ... }
    if (strncmp(line, "dist ", 5) == 0) {
        char name[64], rest[256];
        if (sscanf(line+5, "%63s = %255[^\n]", name, rest) == 2) {
            Tensor t; init_tensor(&t);
            // check if it's a built-in function call
            char fname[64], args[128];
            if (sscanf(rest, "%63s(%127[^)])", fname, args) == 2) {
                // parse args as comma separated doubles
                double a1=0, a2=0;
                int nargs = sscanf(args, "%lf, %lf", &a1, &a2);
                if (strcmp(fname, "GAUSSIAN")==0) {
                    if (nargs==2) t = gaussian(a1, a2);
                    else t = gaussian(0,1);
                } else if (strcmp(fname, "UNIFORM")==0) {
                    if (nargs==2) t = uniform(a1, a2);
                    else t = uniform(0,1);
                } else if (strcmp(fname, "BINOMIAL")==0) {
                    if (nargs==2) t = binomial((int)a1, a2);
                    else t = binomial(1,0.5);
                } else if (strcmp(fname, "EXPONENTIAL")==0) {
                    t = exponential(a1);
                } else if (strcmp(fname, "POISSON")==0) {
                    t = poisson(a1);
                } else if (strcmp(fname, "LOGNORMAL")==0) {
                    t = lognormal(a1, a2);
                } else {
                    printf("Unknown distribution: %s\n", fname);
                    return;
                }
            } else if (rest[0] == '{') {
                // explicit states
                char* p = rest+1;
                while (*p && *p != '}') {
                    double val, prob;
                    int n=0;
                    if (sscanf(p, "%lf : %lf %n", &val, &prob, &n) == 2) {
                        add_state(&t, val, prob);
                        p += n;
                        while (isspace(*p)) p++;
                        if (*p == ',') p++;
                        while (isspace(*p)) p++;
                    } else break;
                }
                normalize(&t);
            } else {
                // dist name = other_tensor
                char other[64];
                sscanf(rest, "%63s", other);
                Symbol* s = find_sym(other);
                if (s && s->type == SYM_TENSOR) t = s->value.t;
                else { printf("Undefined tensor: %s\n", other); return; }
            }
            // remove existing symbol of same name
            for (int i=0; i<sym_count; i++) {
                if (strcmp(symtab[i].name, name)==0) {
                    // shift left
                    for (int j=i; j<sym_count-1; j++) symtab[j]=symtab[j+1];
                    sym_count--;
                    break;
                }
            }
            add_tensor(name, &t);
            printf("Defined tensor '%s'\n", name);
        }
        return;
    }
    
    // double var = measure dist  or  double var = mean(dist)  etc.
    if (strncmp(line, "double ", 7) == 0) {
        char name[64], rest[256];
        if (sscanf(line+7, "%63s = %255[^\n]", name, rest) == 2) {
            char* p = rest;
            while (isspace(*p)) p++;
            if (strncmp(p, "measure", 7) == 0) {
                p += 7; while(isspace(*p)) p++;
                char tname[64];
                sscanf(p, "%63s", tname);
                Symbol* s = find_sym(tname);
                if (s && s->type == SYM_TENSOR) {
                    double val = measure(&s->value.t);
                    add_scalar(name, val);
                    printf("%s = %.6f\n", name, val);
                } else printf("Tensor %s not found\n", tname);
            } else if (strncmp(p, "mean", 4) == 0) {
                p += 4; while(isspace(*p)) p++;
                if (*p=='(') p++;
                char tname[64];
                sscanf(p, "%63s", tname);
                Symbol* s = find_sym(tname);
                if (s && s->type == SYM_TENSOR) {
                    double val = mean(&s->value.t);
                    add_scalar(name, val);
                    printf("%s = %.6f\n", name, val);
                } else printf("Tensor %s not found\n", tname);
            } else if (strncmp(p, "variance", 8) == 0) {
                p += 8; while(isspace(*p)) p++;
                if (*p=='(') p++;
                char tname[64];
                sscanf(p, "%63s", tname);
                Symbol* s = find_sym(tname);
                if (s && s->type == SYM_TENSOR) {
                    double val = variance(&s->value.t);
                    add_scalar(name, val);
                    printf("%s = %.6f\n", name, val);
                } else printf("Tensor %s not found\n", tname);
            } else if (strncmp(p, "quantile", 8) == 0) {
                p += 8; while(isspace(*p)) p++;
                if (*p=='(') p++;
                char tname[64];
                double q;
                sscanf(p, "%63s , %lf", tname, &q);
                Symbol* s = find_sym(tname);
                if (s && s->type == SYM_TENSOR) {
                    double val = quantile(&s->value.t, q);
                    add_scalar(name, val);
                    printf("%s = %.6f\n", name, val);
                } else printf("Tensor %s not found\n", tname);
            } else {
                // expression assignment
                expr_str = p;
                expr_pos = 0;
                double val = expr_eval();
                add_scalar(name, val);
                printf("%s = %.6f\n", name, val);
            }
        }
        return;
    }
    
    // print expr
    if (strncmp(line, "print", 5) == 0) {
        char* p = line+5; while(isspace(*p)) p++;
        if (*p == '"') {
            p++; char* end = strchr(p, '"');
            if (end) *end = 0;
            printf("%s\n", p);
        } else {
            char name[64];
            sscanf(p, "%63s", name);
            Symbol* s = find_sym(name);
            if (s) {
                if (s->type == SYM_TENSOR) print_tensor(&s->value.t);
                else printf("%.6f\n", s->value.d);
            } else printf("Undefined symbol: %s\n", name);
        }
        return;
    }
    
    // clear
    if (strcmp(line, "clear") == 0) { printf("\033[2J\033[H"); return; }
    
    // color N
    if (strncmp(line, "color", 5) == 0) {
        int c; sscanf(line+5, "%d", &c);
        printf("\033[3%dm", c%8);
        return;
    }
    
    // getc var
    if (strncmp(line, "getc", 4) == 0) {
        char name[64];
        sscanf(line+4, "%63s", name);
        int ch = getchar();
        add_scalar(name, (double)ch);
        printf("%s = %d ('%c')\n", name, ch, ch);
        return;
    }
    
    // putc var
    if (strncmp(line, "putc", 4) == 0) {
        char name[64];
        sscanf(line+4, "%63s", name);
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_SCALAR)
            putchar((int)s->value.d);
        else printf("Undefined scalar: %s\n", name);
        return;
    }
    
    // exit
    if (strncmp(line, "exit", 4) == 0) {
        exit(0);
    }
    
    // measure dist (immediate)
    if (strncmp(line, "measure", 7) == 0) {
        char* p = line+7; while(isspace(*p)) p++;
        char name[64];
        sscanf(p, "%63s", name);
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_TENSOR) {
            double val = measure(&s->value.t);
            printf("measure(%s) = %.6f\n", name, val);
        } else printf("Tensor %s not found\n", name);
        return;
    }
    
    // mean(dist)
    if (strncmp(line, "mean", 4) == 0) {
        char* p = line+4; while(isspace(*p)) p++;
        if (*p=='(') p++;
        char name[64];
        sscanf(p, "%63s", name);
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_TENSOR) {
            printf("mean(%s) = %.6f\n", name, mean(&s->value.t));
        } else printf("Tensor %s not found\n", name);
        return;
    }
    
    // variance(dist)
    if (strncmp(line, "variance", 8) == 0) {
        char* p = line+8; while(isspace(*p)) p++;
        if (*p=='(') p++;
        char name[64];
        sscanf(p, "%63s", name);
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_TENSOR) {
            printf("variance(%s) = %.6f\n", name, variance(&s->value.t));
        } else printf("Tensor %s not found\n", name);
        return;
    }
    
    // quantile(dist, p)
    if (strncmp(line, "quantile", 8) == 0) {
        char* p = line+8; while(isspace(*p)) p++;
        if (*p=='(') p++;
        char name[64]; double q;
        sscanf(p, "%63s , %lf", name, &q);
        Symbol* s = find_sym(name);
        if (s && s->type == SYM_TENSOR) {
            printf("quantile(%s, %.4f) = %.6f\n", name, q, quantile(&s->value.t, q));
        } else printf("Tensor %s not found\n", name);
        return;
    }
    
    // add dist1 + dist2 (convolution)
    {
        char left[64], right[64];
        if (sscanf(line, "%63s + %63s", left, right) == 2) {
            Symbol* s1 = find_sym(left);
            Symbol* s2 = find_sym(right);
            if (s1 && s1->type==SYM_TENSOR && s2 && s2->type==SYM_TENSOR) {
                Tensor res;
                conv_add(&s1->value.t, &s2->value.t, &res);
                print_tensor(&res);
            } else printf("Both operands must be tensors\n");
            return;
        }
    }
    
    printf("Unrecognized command.\n");
}

// ---------- Utility function ----------
static void trim(char* s) {
    char* end;
    while (isspace((unsigned char)*s)) s++;
    if (*s==0) return;
    end = s + strlen(s) - 1;
    while (end > s && isspace((unsigned char)*end)) end--;
    end[1] = '\0';
}

// ---------- REPL Main ----------
static void repl(void) {
    char line[1024];
    printf("PROBOL Interactive REPL (Standard Library v1.0)\n");
    printf("Type 'exit' to quit.\n");
    while (1) {
        printf("probol> ");
        fflush(stdout);
        if (!fgets(line, sizeof(line), stdin)) break;
        process_line(line);
    }
}

int main(int argc, char** argv) {
    // Preload some standard scalars
    add_scalar("pi", 3.141592653589793);
    add_scalar("e", 2.718281828459045);
    repl();
    return 0;
}