#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <random>
#include <cmath>
#include <algorithm>
#include <chrono>
#include <unordered_map>
#include <numeric>
#include <iomanip>

// ---------------------------------------------------------------------------
// JSON parsing helpers (minimal)
// ---------------------------------------------------------------------------

std::string json_unescape(const std::string& s) {
    std::string out;
    for (size_t i = 0; i < s.size(); i++) {
        if (s[i] == '\\' && i + 1 < s.size()) {
            char next = s[i+1];
            if (next == 'n') out += '\n';
            else if (next == 't') out += '\t';
            else if (next == 'r') out += '\r';
            else out += next;
            i++;
        } else if (s[i] != '"') {
            out += s[i];
        }
    }
    return out;
}

std::string extract_json_string(const std::string& line, const std::string& key) {
    size_t pos = line.find("\"" + key + "\":");
    if (pos == std::string::npos) return "";
    pos = line.find("\"", pos + key.size() + 3);
    if (pos == std::string::npos) return "";
    size_t end = pos + 1;
    while (end < line.size() && (line[end] != '"' || line[end-1] == '\\')) end++;
    return json_unescape(line.substr(pos + 1, end - pos - 1));
}

// ---------------------------------------------------------------------------
// Feature extraction
// ---------------------------------------------------------------------------

const std::vector<std::string> local_areas = {
    "manhattan", "chelsea", "midtown", "hell's kitchen", "upper west side",
    "upper east side", "tribeca", "soho", "west village", "east village",
    "financial district", "flatiron", "gramercy", "murray hill", "nomad", "nyc", "new york"
};

const std::vector<std::string> service_words = {
    "deep tissue", "sports", "recovery", "massage", "bodywork", "relief",
    "therapy", "knots", "tension", "pressure", "shoulder", "back", "hip",
    "neck", "glute", "stress", "desk", "travel", "athlete", "professional"
};

const std::vector<std::string> cta_words = {
    "message", "text", "email", "book", "call", "contact", "dm", "schedule", "reach out"
};

const std::vector<std::string> trust_words = {
    "professional", "clean", "private", "respect", "discreet", "boundaries", "communication"
};

const std::vector<std::string> humor_words = {
    "wolf", "robot", "concrete", "phone cord", "group chat", "screaming", "dramatic", "hostile", "no fluff", "feather", "magic", "gps"
};

std::string to_lower(const std::string& s) {
    std::string r = s;
    std::transform(r.begin(), r.end(), r.begin(), ::tolower);
    return r;
}

std::vector<std::string> tokenize(const std::string& text) {
    std::vector<std::string> tokens;
    std::string cur;
    for (char c : text) {
        if (std::isalpha(static_cast<unsigned char>(c)) || c == '\'') {
            cur += std::tolower(c);
        } else if (!cur.empty()) {
            tokens.push_back(cur);
            cur.clear();
        }
    }
    if (!cur.empty()) tokens.push_back(cur);
    return tokens;
}

std::vector<std::string> split_sentences(const std::string& text) {
    std::vector<std::string> sents;
    std::string cur;
    for (char c : text) {
        cur += c;
        if (c == '.' || c == '!' || c == '?') {
            if (!cur.empty()) {
                sents.push_back(cur);
                cur.clear();
            }
        }
    }
    if (!cur.empty()) sents.push_back(cur);
    return sents;
}

int syllable_count(const std::string& word) {
    std::string w = to_lower(word);
    std::string vowels = "aeiouy";
    int count = 0;
    bool prev_vowel = false;
    for (char c : w) {
        bool is_vowel = vowels.find(c) != std::string::npos;
        if (is_vowel && !prev_vowel) count++;
        prev_vowel = is_vowel;
    }
    if (w.size() > 1 && w.back() == 'e') count--;
    return std::max(1, count);
}

std::vector<double> extract_features(const std::string& headline, const std::string& description) {
    std::string text = to_lower(headline + " " + description);
    auto words = tokenize(text);
    auto sentences = split_sentences(description);

    double headline_len = headline.size() / 100.0;
    double desc_len = description.size() / 500.0;
    double word_count = words.size() / 100.0;
    double avg_sentence_len = 0;
    for (const auto& s : sentences) avg_sentence_len += tokenize(s).size();
    avg_sentence_len = (avg_sentence_len / std::max(1, (int)sentences.size())) / 30.0;
    double paragraph_count = std::min(5.0, (double)std::count(description.begin(), description.end(), '\n') + 1) / 5.0;

    double local_score = 0, service_score = 0, cta_score = 0, trust_score = 0, humor_score = 0;
    for (const auto& a : local_areas) if (text.find(a) != std::string::npos) local_score += 1.0;
    for (const auto& a : service_words) if (text.find(a) != std::string::npos) service_score += 1.0;
    for (const auto& a : cta_words) if (text.find(a) != std::string::npos) cta_score += 1.0;
    for (const auto& a : trust_words) if (text.find(a) != std::string::npos) trust_score += 1.0;
    for (const auto& a : humor_words) if (text.find(a) != std::string::npos) humor_score += 1.0;
    local_score /= local_areas.size();
    service_score /= service_words.size();
    cta_score /= cta_words.size();
    trust_score /= trust_words.size();
    humor_score = std::min(1.0, humor_score / 3.0);

    // Readability
    double readability;
    if (description.size() >= 50 && description.size() <= 300) readability = 1.0;
    else if (description.size() < 50) readability = 0.5;
    else readability = std::max(0.3, 1.0 - (description.size() - 300.0) / 500.0);

    // Speech score
    double avg_word_len = 0, avg_syllables = 0;
    for (const auto& w : words) { avg_word_len += w.size(); avg_syllables += syllable_count(w); }
    avg_word_len /= std::max(1, (int)words.size());
    avg_syllables /= std::max(1, (int)words.size());
    double simple_ratio = 0;
    for (const auto& w : words) if (w.size() <= 6 && w.find('\'') == std::string::npos) simple_ratio += 1.0;
    simple_ratio /= std::max(1, (int)words.size());
    double speech_score = (1 - std::min(avg_word_len / 8.0, 1.0)) * 0.3
                        + (1 - std::min(avg_syllables / 2.5, 1.0)) * 0.3
                        + (1 - std::min(avg_sentence_len * 30.0 / 20.0, 1.0)) * 0.25
                        + simple_ratio * 0.15;

    return {headline_len, desc_len, word_count, avg_sentence_len, paragraph_count,
            (double)std::count(text.begin(), text.end(), '?') / 2.0,
            (double)std::count(text.begin(), text.end(), '!') / 2.0,
            local_score, service_score, cta_score, trust_score, humor_score, readability, speech_score};
}

int feature_count() { return 14; }

// ---------------------------------------------------------------------------
// Heuristic scoring (to produce synthetic labels)
// ---------------------------------------------------------------------------

std::vector<double> heuristic_score(const std::string& headline, const std::string& description) {
    std::string text = to_lower(headline + " " + description);
    double risk = 0;
    if (text.find("cure") != std::string::npos || text.find("medical") != std::string::npos || text.find("heal") != std::string::npos) risk += 0.4;
    if (text.find("guarantee") != std::string::npos || text.find("100%") != std::string::npos) risk += 0.5;
    if (text.find("sex") != std::string::npos || text.find("escort") != std::string::npos) risk += 1.0;
    if (text.find("fake") != std::string::npos) risk += 1.0;
    double pos = 0, neg = 0;
    std::vector<std::string> pos_words = {"better", "relief", "recover", "strong", "calm", "clean", "professional", "help", "good", "best", "real", "focus"};
    std::vector<std::string> neg_words = {"hurt", "pain", "stress", "tight", "bad", "dramatic", "tension", "stiff", "tired"};
    for (const auto& w : pos_words) if (text.find(w) != std::string::npos) pos++;
    for (const auto& w : neg_words) if (text.find(w) != std::string::npos) neg++;
    double sentiment = (pos + neg) > 0 ? pos / (pos + neg) : 0.5;
    double cta = (text.find("message") != std::string::npos || text.find("text") != std::string::npos || text.find("email") != std::string::npos || text.find("book") != std::string::npos) ? 1.0 : 0.0;
    double local = (text.find("manhattan") != std::string::npos) ? 1.0 : 0.0;
    double clarity = (description.size() >= 50 && description.size() <= 300) ? 1.0 : 0.7;
    double composite = (1 - risk) * 0.35 + sentiment * 0.25 + cta * 0.15 + local * 0.15 + clarity * 0.1;
    return {
        std::min(0.08, composite * 0.08 + 0.01),
        std::min(0.05, composite * 0.05 + 0.005),
        std::min(0.03, composite * 0.03 + 0.003)
    };
}

// ---------------------------------------------------------------------------
// MLP with backpropagation
// ---------------------------------------------------------------------------

struct MLP {
    int input_size, hidden_size, output_size;
    std::vector<std::vector<double>> W1, W2;
    std::vector<double> b1, b2;
    std::vector<std::vector<double>> mW1, vW1, mW2, vW2;
    std::vector<double> mb1, vb1, mb2, vb2;
    int t = 0;
    double dropout;

    MLP(int in_s, int hid_s, int out_s, double seed, double drop = 0.2)
        : input_size(in_s), hidden_size(hid_s), output_size(out_s), dropout(drop) {
        std::mt19937 rng((unsigned)seed);
        std::normal_distribution<double> dist(0.0, 1.0);
        W1.resize(input_size, std::vector<double>(hidden_size));
        mW1.resize(input_size, std::vector<double>(hidden_size, 0.0));
        vW1.resize(input_size, std::vector<double>(hidden_size, 0.0));
        for (int i = 0; i < input_size; i++)
            for (int j = 0; j < hidden_size; j++)
                W1[i][j] = dist(rng) * std::sqrt(2.0 / input_size);
        b1.resize(hidden_size, 0.0);
        mb1.resize(hidden_size, 0.0);
        vb1.resize(hidden_size, 0.0);
        W2.resize(hidden_size, std::vector<double>(output_size));
        mW2.resize(hidden_size, std::vector<double>(output_size, 0.0));
        vW2.resize(hidden_size, std::vector<double>(output_size, 0.0));
        for (int i = 0; i < hidden_size; i++)
            for (int j = 0; j < output_size; j++)
                W2[i][j] = dist(rng) * std::sqrt(2.0 / hidden_size);
        b2.resize(output_size, 0.0);
        mb2.resize(output_size, 0.0);
        vb2.resize(output_size, 0.0);
    }

    double relu(double x) { return std::max(0.0, x); }
    double sigmoid(double x) { return 1.0 / (1.0 + std::exp(-std::min(500.0, std::max(-500.0, x)))); }

    std::vector<double> forward(const std::vector<double>& x, bool training, std::vector<double>& z1_out, std::vector<double>& a1_out, std::vector<double>& mask) {
        z1_out.resize(hidden_size);
        a1_out.resize(hidden_size);
        mask.resize(hidden_size);
        std::mt19937 rng((unsigned)std::chrono::steady_clock::now().time_since_epoch().count());
        for (int j = 0; j < hidden_size; j++) {
            z1_out[j] = b1[j];
            for (int i = 0; i < input_size; i++) z1_out[j] += x[i] * W1[i][j];
            a1_out[j] = relu(z1_out[j]);
            if (training) {
                mask[j] = (rng() / (double)rng.max() > dropout) ? 1.0 / (1.0 - dropout) : 0.0;
                a1_out[j] *= mask[j];
            }
        }
        std::vector<double> z2(output_size);
        for (int j = 0; j < output_size; j++) {
            z2[j] = b2[j];
            for (int i = 0; i < hidden_size; i++) z2[j] += a1_out[i] * W2[i][j];
        }
        std::vector<double> a2(output_size);
        for (int j = 0; j < output_size; j++) a2[j] = sigmoid(z2[j]);
        return a2;
    }

    void backward(const std::vector<double>& x, const std::vector<double>& y, const std::vector<double>& y_pred,
                  const std::vector<double>& z1, const std::vector<double>& a1, const std::vector<double>& mask,
                  double lr, double beta1 = 0.9, double beta2 = 0.999, double eps = 1e-8) {
        t++;
        int m = 1;
        std::vector<double> dz2(output_size);
        for (int j = 0; j < output_size; j++) dz2[j] = (y_pred[j] - y[j]) / m;
        std::vector<std::vector<double>> dW2(hidden_size, std::vector<double>(output_size, 0.0));
        std::vector<double> db2(output_size, 0.0);
        for (int j = 0; j < output_size; j++) {
            for (int i = 0; i < hidden_size; i++) dW2[i][j] += a1[i] * dz2[j];
            db2[j] += dz2[j];
        }
        std::vector<double> da1(hidden_size, 0.0);
        for (int i = 0; i < hidden_size; i++) {
            for (int j = 0; j < output_size; j++) da1[i] += dz2[j] * W2[i][j];
            da1[i] *= mask[i];
        }
        std::vector<std::vector<double>> dW1(input_size, std::vector<double>(hidden_size, 0.0));
        std::vector<double> db1(hidden_size, 0.0);
        for (int j = 0; j < hidden_size; j++) {
            double dz1 = da1[j] * (z1[j] > 0 ? 1.0 : 0.0);
            for (int i = 0; i < input_size; i++) dW1[i][j] += x[i] * dz1;
            db1[j] += dz1;
        }

        auto adam = [&](double& W, double& mW, double& vW, double dW) {
            mW = beta1 * mW + (1 - beta1) * dW;
            vW = beta2 * vW + (1 - beta2) * dW * dW;
            double m_hat = mW / (1 - std::pow(beta1, t));
            double v_hat = vW / (1 - std::pow(beta2, t));
            W -= lr * m_hat / (std::sqrt(v_hat) + eps);
        };
        for (int i = 0; i < input_size; i++)
            for (int j = 0; j < hidden_size; j++)
                adam(W1[i][j], mW1[i][j], vW1[i][j], dW1[i][j]);
        for (int j = 0; j < hidden_size; j++) adam(b1[j], mb1[j], vb1[j], db1[j]);
        for (int i = 0; i < hidden_size; i++)
            for (int j = 0; j < output_size; j++)
                adam(W2[i][j], mW2[i][j], vW2[i][j], dW2[i][j]);
        for (int j = 0; j < output_size; j++) adam(b2[j], mb2[j], vb2[j], db2[j]);
    }

    void train_epoch(const std::vector<std::vector<double>>& X, const std::vector<std::vector<double>>& y,
                     double lr, int batch_size, std::mt19937& rng) {
        int n = X.size();
        std::vector<int> idx(n);
        std::iota(idx.begin(), idx.end(), 0);
        std::shuffle(idx.begin(), idx.end(), rng);
        for (int i = 0; i < n; i += batch_size) {
            for (int k = i; k < std::min(n, i + batch_size); k++) {
                int ii = idx[k];
                std::vector<double> z1, a1, mask;
                auto pred = forward(X[ii], true, z1, a1, mask);
                backward(X[ii], y[ii], pred, z1, a1, mask, lr);
            }
        }
    }

    std::vector<double> predict(const std::vector<double>& x) {
        std::vector<double> z1, a1, mask;
        return forward(x, false, z1, a1, mask);
    }
};

// ---------------------------------------------------------------------------
// Metrics
// ---------------------------------------------------------------------------

struct Metrics {
    double mae, rmse, r2;
};

Metrics evaluate(MLP& model, const std::vector<std::vector<double>>& X, const std::vector<std::vector<double>>& y) {
    int n = X.size();
    double mae = 0, mse = 0;
    std::vector<double> y_mean(3, 0);
    for (const auto& row : y) for (int j = 0; j < 3; j++) y_mean[j] += row[j];
    for (double& v : y_mean) v /= n;
    double ss_tot = 0;
    for (int i = 0; i < n; i++) {
        auto pred = model.predict(X[i]);
        for (int j = 0; j < 3; j++) {
            double diff = pred[j] - y[i][j];
            mae += std::abs(diff);
            mse += diff * diff;
            ss_tot += (y[i][j] - y_mean[j]) * (y[i][j] - y_mean[j]);
        }
    }
    mae /= (n * 3);
    mse /= (n * 3);
    return {mae, std::sqrt(mse), 1 - mse * (n * 3) / ss_tot};
}

// ---------------------------------------------------------------------------
// Data loading
// ---------------------------------------------------------------------------

struct Bio {
    std::string headline, description;
    std::vector<double> features;
    std::vector<double> labels;
};

std::vector<Bio> load_bios(const std::string& path, int limit = 0) {
    std::vector<Bio> bios;
    std::ifstream f(path);
    std::string line;
    while (std::getline(f, line)) {
        if (line.empty()) continue;
        Bio b;
        b.headline = extract_json_string(line, "headline");
        b.description = extract_json_string(line, "description");
        b.features = extract_features(b.headline, b.description);
        b.labels = heuristic_score(b.headline, b.description);
        bios.push_back(b);
        if (limit > 0 && (int)bios.size() >= limit) break;
    }
    return bios;
}

// ---------------------------------------------------------------------------
// k-fold cross-validation
// ---------------------------------------------------------------------------

void k_fold_cv(std::vector<Bio>& bios, int k, int epochs, double lr, int batch_size) {
    int n = bios.size();
    int fold_size = n / k;
    std::vector<Metrics> metrics;
    for (int fold = 0; fold < k; fold++) {
        int val_start = fold * fold_size;
        int val_end = val_start + fold_size;
        std::vector<std::vector<double>> X_train, y_train, X_val, y_val;
        for (int i = 0; i < n; i++) {
            if (i >= val_start && i < val_end) {
                X_val.push_back(bios[i].features);
                y_val.push_back(bios[i].labels);
            } else {
                X_train.push_back(bios[i].features);
                y_train.push_back(bios[i].labels);
            }
        }
        MLP model(feature_count(), 32, 3, fold + 1);
        std::mt19937 rng(fold + 1);
        for (int e = 0; e < epochs; e++) model.train_epoch(X_train, y_train, lr, batch_size, rng);
        Metrics m = evaluate(model, X_val, y_val);
        metrics.push_back(m);
        std::cout << "Fold " << fold << " — MAE=" << m.mae << " RMSE=" << m.rmse << " R2=" << m.r2 << std::endl;
    }
    double mae = 0, rmse = 0, r2 = 0;
    for (const auto& m : metrics) { mae += m.mae; rmse += m.rmse; r2 += m.r2; }
    std::cout << "CV mean — MAE=" << mae / k << " RMSE=" << rmse / k << " R2=" << r2 / k << std::endl;
}

// ---------------------------------------------------------------------------
// Walk-forward validation
// ---------------------------------------------------------------------------

void walk_forward(std::vector<Bio>& bios, double train_ratio, int epochs, double lr, int batch_size) {
    int n = bios.size();
    int split = (int)(train_ratio * n);
    std::vector<std::vector<double>> X_train, y_train, X_val, y_val;
    for (int i = 0; i < n; i++) {
        if (i < split) { X_train.push_back(bios[i].features); y_train.push_back(bios[i].labels); }
        else { X_val.push_back(bios[i].features); y_val.push_back(bios[i].labels); }
    }
    MLP model(feature_count(), 32, 3, 42);
    std::mt19937 rng(42);
    for (int e = 0; e < epochs; e++) {
        model.train_epoch(X_train, y_train, lr, batch_size, rng);
        if (e % 20 == 0) {
            Metrics m = evaluate(model, X_train, y_train);
            std::cout << "Epoch " << e << " train_loss=" << m.mae << std::endl;
        }
    }
    Metrics m = evaluate(model, X_val, y_val);
    std::cout << "Walk-forward — MAE=" << m.mae << " RMSE=" << m.rmse << " R2=" << m.r2 << std::endl;
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cout << "Usage: bio_ml <bios.jsonl> [cv|walk|train] [limit] [epochs] [lr] [batch_size]" << std::endl;
        return 1;
    }
    std::string path = argv[1];
    std::string mode = argc > 2 ? argv[2] : "cv";
    int limit = argc > 3 ? std::stoi(argv[3]) : 10000;
    int epochs = argc > 4 ? std::stoi(argv[4]) : 100;
    double lr = argc > 5 ? std::stod(argv[5]) : 0.01;
    int batch_size = argc > 6 ? std::stoi(argv[6]) : 64;

    auto start = std::chrono::high_resolution_clock::now();
    std::cout << "Loading bios..." << std::endl;
    auto bios = load_bios(path, limit);
    std::cout << "Loaded " << bios.size() << " bios" << std::endl;

    if (mode == "cv") k_fold_cv(bios, 5, epochs, lr, batch_size);
    else if (mode == "walk") walk_forward(bios, 0.8, epochs, lr, batch_size);

    auto end = std::chrono::high_resolution_clock::now();
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << "Total time: " << ms << " ms" << std::endl;
    return 0;
}