// maxent.js BAYES STAR (c) coppola.ai
const mongoose = require("mongoose")
const models = require("./models")

const weightSchema = new mongoose.Schema({
    feature: { type: String, required: true },
    weight: { type: Number, required: true },
})
const WeightRecord = mongoose.model('Weight', weightSchema);

function RandomWeight() {
    return (Math.random() - Math.random()) / 5
}

function PositiveFeature(feature) {
    return '+++' + feature + '+++'
}

function NegativeFeature(feature) {
    return '---' + feature + '---'
}

async function InitializeWeights(implication) {
    const weight = RandomWeight()
    const feature = implication.UniqueKey()
    const pwr = WeightRecord({ feature: PositiveFeature(feature), weight })
    await pwr.save()
    const nwr = WeightRecord({ feature: NegativeFeature(feature), weight })
    await nwr.save()
}

async function GetPropositionProbability(searchString) {
    const record = await models.PropositionRecord.find({ searchString })
    const r = record[0].probability
    return r
}

function Sigmoid(x) {
    return 1 / (1 + Math.exp(-x));
}

function DotProduct(dict1, dict2) {
    let result = 0;
    for (const key of Object.keys(dict1)) {
        const v1 = dict1[key]
        const v2 = dict2[key]
        if (v1 == null || v2 == null) {
        } else {
            result += v1 * v2
        }
    }
    return result;
}

async function FeaturesFromBacklinks(backlinks) {
    var result = {}
    for (var i = 0; i < backlinks.length; i++) {
        const backlink = backlinks[i]
        const feature = backlink.implication.UniqueKey()
        const searchString = backlink.proposition.SearchString()
        const probability = await GetPropositionProbability(searchString)
        result[PositiveFeature(feature)] = probability
        result[NegativeFeature(feature)] = 1 - probability
    }
    return result
}

function ComputeProbability(weights, features) {
    const dot = DotProduct(weights, features);
    const probability = Sigmoid(dot);
    return probability;
}

async function ReadWeights(features) {
    var r = {}
    for (const feature of Object.keys(features)) {
        const record = await WeightRecord.findOne({ feature })
        if (record) {
            r[feature] = record.weight
        }
    }
    return r
}

async function SaveWeights(weights) {
    for (const feature of Object.keys(weights)) {
        const updated = await WeightRecord.findOneAndUpdate({ feature }, { weight: weights[feature] }, { new: true, runValidators: true });
    }
}

function ComputeExpectedFeatures(probability, features) {
    let r = {};
    for (let key in features) {
        r[key] = features[key] * probability;
    }
    return r;
}

const LEARNING_RATE = 0.1

function DoSGDUpdate(weights, goldFeatures, expectedFeatures) {
    var r = {}
    for (const feature of Object.keys(weights)) {
        const wv = weights[feature]
        const gv = goldFeatures[feature]
        const ev = expectedFeatures[feature]
        const newWeight = wv + LEARNING_RATE * (gv - ev)
        console.log({ feature, wv, gv, ev, newWeight }, DoSGDUpdate)
        r[feature] = newWeight
    }
    return r
}

async function TrainOnExample(proposition, backlinks) {
    const features = await FeaturesFromBacklinks(backlinks)
    const weightVector = await ReadWeights(features)
    const probability = ComputeProbability(weightVector, features);
    const expected = ComputeExpectedFeatures(probability, features)
    const newWeight = DoSGDUpdate(weightVector, features, expected)
    await SaveWeights(newWeight)
}

async function DumpWeights() {
    const records = await WeightRecord.find({})
    console.log({ records }, DumpWeights)
}

module.exports = { InitializeWeights, TrainOnExample, DumpWeights }