#include "probol_stdlib.h"
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>
#include <termios.h>

/*----------------------------------------------------------------------------
  Internal helpers
----------------------------------------------------------------------------*/
static int compar_prob_pair(const void *a, const void *b) {
    const prob_fixed_t *va = a;
    const prob_fixed_t *vb = b;
    return (va < vb) ? -1 : (va > vb);
}

static void sort_tensor_by_value(ProbTensor *t) {
    /* simple insertion sort – small sizes */
    for (int i = 1; i < t->count; i++) {
        prob_fixed_t v = t->values[i];
        prob_fixed_t p = t->probs[i];
        int j = i - 1;
        while (j >= 0 && t->values[j] > v) {
            t->values[j+1] = t->values[j];
            t->probs[j+1] = t->probs[j];
            j--;
        }
        t->values[j+1] = v;
        t->probs[j+1] = p;
    }
}

/*----------------------------------------------------------------------------
  Probability Tensor Primitives
----------------------------------------------------------------------------*/
void probol_init_tensor(ProbTensor *t) {
    t->count = 0;
    memset(t->values, 0, sizeof(t->values));
    memset(t->probs, 0, sizeof(t->probs));
}

void probol_add_state(ProbTensor *t, prob_fixed_t val, prob_fixed_t prob) {
    if (t->count >= PROBOL_MAX_STATES) return;
    /* merge if value already exists */
    for (int i = 0; i < t->count; i++) {
        if (t->values[i] == val) {
            t->probs[i] = fp_add(t->probs[i], prob);
            return;
        }
    }
    t->values[t->count] = val;
    t->probs[t->count] = prob;
    t->count++;
}

void probol_normalize(ProbTensor *t) {
    prob_fixed_t sum = 0;
    for (int i = 0; i < t->count; i++) sum = fp_add(sum, t->probs[i]);
    if (sum == 0) return;
    for (int i = 0; i < t->count; i++) t->probs[i] = fp_div(t->probs[i], sum);
}

void probol_add_tensors(const ProbTensor *A, const ProbTensor *B, ProbTensor *C) {
    probol_init_tensor(C);
    for (int i = 0; i < A->count; i++) {
        for (int j = 0; j < B->count; j++) {
            prob_fixed_t val = fp_add(A->values[i], B->values[j]);
            prob_fixed_t prob = fp_mul(A->probs[i], B->probs[j]);
            probol_add_state(C, val, prob);
        }
    }
    probol_normalize(C);
}

void probol_multiply_tensors(const ProbTensor *A, const ProbTensor *B, ProbTensor *C) {
    probol_init_tensor(C);
    for (int i = 0; i < A->count; i++) {
        for (int j = 0; j < B->count; j++) {
            prob_fixed_t val = fp_mul(A->values[i], B->values[j]);
            prob_fixed_t prob = fp_mul(A->probs[i], B->probs[j]);
            probol_add_state(C, val, prob);
        }
    }
    probol_normalize(C);
}

prob_fixed_t probol_quantile(const ProbTensor *t, prob_fixed_t quantile) {
    ProbTensor tmp;
    memcpy(&tmp, t, sizeof(ProbTensor));
    sort_tensor_by_value(&tmp);
    prob_fixed_t cum = 0;
    for (int i = 0; i < tmp.count; i++) {
        cum = fp_add(cum, tmp.probs[i]);
        if (cum >= quantile) return tmp.values[i];
    }
    return (tmp.count ? tmp.values[tmp.count-1] : 0);
}

prob_fixed_t probol_measure(const ProbTensor *t) {
    prob_fixed_t r = (prob_fixed_t)(rand() % FP_ONE);
    prob_fixed_t cum = 0;
    for (int i = 0; i < t->count; i++) {
        cum = fp_add(cum, t->probs[i]);
        if (r < cum) return t->values[i];
    }
    return (t->count ? t->values[t->count-1] : 0);
}

prob_fixed_t probol_prob_condition(const ProbTensor *t, bool (*condition)(prob_fixed_t)) {
    prob_fixed_t p = 0;
    for (int i = 0; i < t->count; i++) {
        if (condition(t->values[i])) p = fp_add(p, t->probs[i]);
    }
    return p;
}

prob_fixed_t probol_entropy(const ProbTensor *t) {
    prob_fixed_t h = 0;
    for (int i = 0; i < t->count; i++) {
        prob_fixed_t p = t->probs[i];
        if (p != 0) {
            /* H = - Σ p * log2(p). Use fixed‑point approximate log */
            double pd = FP_TO_DOUBLE(p);
            h = fp_sub(h, fp_mul(p, FP_FROM(log2(pd))));
        }
    }
    return h;
}

/*----------------------------------------------------------------------------
  Parametric Distributions (Normal)
----------------------------------------------------------------------------*/
void probol_gaussian(prob_fixed_t mu, prob_fixed_t sigma, ProbTensor *out) {
    probol_init_tensor(out);
    /* crude discretization: 7 points symmetric around mu */
    for (int i = -3; i <= 3; i++) {
        prob_fixed_t v = fp_add(mu, fp_mul(FP_FROM(i), sigma));
        double z = i; /* standard deviation multiplier */
        double prob = exp(-z*z/2.0) / sqrt(2.0 * M_PI);
        probol_add_state(out, v, FP_FROM(prob));
    }
    probol_normalize(out);
}

/*----------------------------------------------------------------------------
  Copula Implementation (simplified but functional)
----------------------------------------------------------------------------*/
void probol_copula_init_gaussian(CopulaGaussian *c, int dim) {
    c->dim = dim;
    for (int i = 0; i < dim; i++)
        for (int j = 0; j < dim; j++)
            c->corr[i][j] = (i == j) ? FP_ONE : 0;
}

void probol_copula_set_correlation(CopulaGaussian *c, int i, int j, prob_fixed_t rho) {
    if (i < c->dim && j < c->dim) {
        c->corr[i][j] = c->corr[j][i] = rho;
    }
}

void probol_copula_compute_cholesky(CopulaGaussian *c) {
    /* Cholesky–Crout algorithm, fixed‑point, assuming positive definite */
    prob_fixed_t *L = &c->cholesky[0][0];
    int n = c->dim;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j <= i; j++) {
            prob_fixed_t sum = 0;
            for (int k = 0; k < j; k++)
                sum = fp_add(sum, fp_mul(c->cholesky[i][k], c->cholesky[j][k]));
            if (i == j)
                c->cholesky[i][j] = (prob_fixed_t)sqrt(FP_TO_DOUBLE(fp_sub(c->corr[i][i], sum)));
            else
                c->cholesky[i][j] = fp_div(fp_sub(c->corr[i][j], sum), c->cholesky[j][j]);
        }
    }
}

void probol_copula_sample(const Copula *c, const prob_fixed_t *uniforms, prob_fixed_t *correlated, int n) {
    /* For Gaussian copula: transform uniforms → normals, multiply by Cholesky */
    if (c->type == COPULA_GAUSSIAN) {
        double u_norm[PROBOL_MAX_ASSETS];
        for (int i = 0; i < n; i++) {
            double u = FP_TO_DOUBLE(uniforms[i]);
            u = fmax(1e-10, fmin(1.0 - 1e-10, u));
            u_norm[i] = sqrt(2.0) * erfinv(2.0 * u - 1.0);
        }
        double res[PROBOL_MAX_ASSETS] = {0};
        for (int i = 0; i < n; i++) {
            for (int j = 0; j <= i; j++) {
                res[i] += FP_TO_DOUBLE(c->gaussian.cholesky[i][j]) * u_norm[j];
            }
        }
        /* convert back to uniform via normal CDF */
        for (int i = 0; i < n; i++) {
            double norm = 0.5 * erfc(-res[i] / M_SQRT2);
            correlated[i] = FP_FROM(norm);
        }
    } else {
        /* stub for other copulas – fallback to independent */
        for (int i = 0; i < n; i++) correlated[i] = uniforms[i];
    }
}

/*----------------------------------------------------------------------------
  GARCH(1,1)
----------------------------------------------------------------------------*/
void probol_garch_update(const GARCH11 *params, prob_fixed_t innovation, prob_fixed_t *variance) {
    prob_fixed_t new_var = fp_add(params->omega,
                         fp_add(fp_mul(params->alpha, fp_mul(innovation, innovation)),
                                fp_mul(params->beta, *variance)));
    *variance = new_var;
}

prob_fixed_t probol_garch_simulate(const GARCH11 *params, int steps, prob_fixed_t *var_out) {
    prob_fixed_t var = params->init_var;
    prob_fixed_t ret = 0;
    for (int t = 0; t < steps; t++) {
        double z = (double)rand() / RAND_MAX * 2.0 - 1.0; /* crude normal */
        prob_fixed_t innov = fp_mul(FP_FROM(z), (prob_fixed_t)sqrt(FP_TO_DOUBLE(var)));
        ret = fp_add(ret, innov);
        probol_garch_update(params, innov, &var);
        if (var_out) var_out[t] = var;
    }
    return ret;
}

/*----------------------------------------------------------------------------
  CCT Scheduler – detects entropy peaks
----------------------------------------------------------------------------*/
static prob_fixed_t dummy_dynamics(int t, const ProbTensor *state) {
    (void)t; (void)state;
    return FP_FROM(0.0); /* placeholder */
}

CCTSchedule probol_cct_schedule(const ProbTensor *initial, int horizon,
                                prob_fixed_t (*dynamics)(int t, const ProbTensor *state)) {
    CCTSchedule sched = { .num_collapses = 0 };
    if (!dynamics) dynamics = dummy_dynamics;
    prob_fixed_t entropy_history[256];
    for (int t = 0; t < horizon && t < 256; t++) {
        ProbTensor state = *initial; /* simplified: propagate with dynamics */
        entropy_history[t] = probol_entropy(&state);
    }
    /* detect spikes where gradient exceeds threshold */
    prob_fixed_t thresh = FP_FROM(0.2);
    for (int t = 1; t < horizon - 1; t++) {
        prob_fixed_t grad1 = fp_sub(entropy_history[t], entropy_history[t-1]);
        prob_fixed_t grad2 = fp_sub(entropy_history[t+1], entropy_history[t]);
        if (grad1 > thresh && grad2 > thresh) {
            sched.collapse_points[sched.num_collapses++] = t;
            if (sched.num_collapses >= 256) break;
        }
    }
    return sched;
}

/*----------------------------------------------------------------------------
  SIMULATEP – time series simulation with copula and optional GARCH
----------------------------------------------------------------------------*/
void probol_simulate(const SimulationConfig *cfg,
                     TimeSeriesState *history,
                     ProbTensor *final_tensor) {
    ProbTensor current = *cfg->marginals; /* start with initial marginals */
    int collapse_idx = 0;
    for (int day = 0; day < cfg->horizon; day++) {
        /* Check if we collapse today (CCT schedule) */
        if (collapse_idx < cfg->collapse_schedule.num_collapses &&
            cfg->collapse_schedule.collapse_points[collapse_idx] == day) {
            /* MEAS: collapse each asset independently */
            prob_fixed_t realized = probol_measure(&current);
            if (history) history[day].value = realized;
            collapse_idx++;
            /* After collapse, rebuild tensor as a single point (deterministic) */
            probol_init_tensor(&current);
            probol_add_state(&current, realized, FP_ONE);
            continue;
        }
        /* Simulate one step using copula and GARCH dynamics */
        prob_fixed_t uniforms[PROBOL_MAX_ASSETS];
        for (int a = 0; a < cfg->num_assets; a++) {
            uniforms[a] = (prob_fixed_t)(rand() % FP_ONE); /* uniform random */
        }
        prob_fixed_t correlated_uniforms[PROBOL_MAX_ASSETS];
        probol_copula_sample(cfg->copula, uniforms, correlated_uniforms, cfg->num_assets);
        /* Update each asset's distribution (simplified: shift by correlated innovation) */
        ProbTensor new_tensor;
        probol_init_tensor(&new_tensor);
        for (int i = 0; i < current.count; i++) {
            for (int a = 0; a < cfg->num_assets; a++) {
                prob_fixed_t innov = fp_sub(correlated_uniforms[a], FP_FROM(0.5));
                prob_fixed_t new_val = fp_add(current.values[i], innov);
                probol_add_state(&new_tensor, new_val, current.probs[i]);
            }
        }
        probol_normalize(&new_tensor);
        current = new_tensor;
        if (history) history[day].value = current.count ? current.values[0] : 0;
    }
    *final_tensor = current;
}

/*----------------------------------------------------------------------------
  I/O Extensions (compatible with probol_compiler)
----------------------------------------------------------------------------*/
void probol_print(const char *str) {
    write(STDOUT_FILENO, str, strlen(str));
}

void probol_print_tensor(const ProbTensor *t) {
    char buf[256];
    int off = snprintf(buf, sizeof(buf), "ProbTensor[%d]:", t->count);
    for (int i = 0; i < t->count && off < (int)sizeof(buf)-32; i++) {
        off += snprintf(buf+off, sizeof(buf)-off, " %ld:%ld",
                        (long)t->values[i], (long)t->probs[i]);
    }
    probol_print(buf);
    probol_print("\n");
}

void probol_clear_screen(void) { probol_print("\033[2J\033[H"); }
void probol_set_color(int code) {
    char buf[16];
    snprintf(buf, sizeof(buf), "\033[3%dm", code & 7);
    probol_print(buf);
}
void probol_gotoxy(int x, int y) {
    char buf[32];
    snprintf(buf, sizeof(buf), "\033[%d;%dH", y, x);
    probol_print(buf);
}
void probol_get_terminal_size(int *width, int *height) {
    struct winsize ws;
    if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0) {
        *width = ws.ws_col;
        *height = ws.ws_row;
    } else {
        *width = 80; *height = 24;
    }
}

static struct termios orig_term;
void probol_raw_mode_enable(void) {
    tcgetattr(STDIN_FILENO, &orig_term);
    struct termios raw = orig_term;
    raw.c_iflag &= ~(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
    raw.c_oflag &= ~OPOST;
    raw.c_cflag |= CS8;
    raw.c_lflag &= ~(ECHO | ICANON | IEXTEN | ISIG);
    raw.c_cc[VMIN] = 0;
    raw.c_cc[VTIME] = 1;
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &raw);
}
void probol_raw_mode_disable(void) {
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &orig_term);
}
char probol_getch(void) {
    char c = 0;
    read(STDIN_FILENO, &c, 1);
    return c;
}