/*
   HARK Library
   Copyright (C) 2021 Honda Research Institute Japan Inc.
 */

#include "hark-lib/node/SemiBlindICA.h"

#include "hark-lib/matrix.h"
#include <fstream>

namespace hark
{
namespace node
{
    SemiBlindICA::SBSemiBlindICA::SBSemiBlindICA(int taps, int itvl, float lambda,
            float mu_s, float mu_v, float mu_n, float is_zero)
    {
        int size;

        // normalization
        mu_s /= 10000;//36768.0;
        mu_v /= 10000;//36768.0;
        mu_n /= 10000;//36768.0;

        //printf("%d %d\n", i, j);
        //id_i = i;
        //id_j = j;

        // set member parameters
        this->itvl = itvl;
        this->taps = taps;
        this->mu_s = mu_s;
        this->mu_n = mu_n;
        this->is_zero = is_zero;

        idx  = 0;
        alpha = 1.0;
        n_var = 1.0;
        r_var = 1.0;
        FLOOR_VALUE = 1.0E-5;

        // filter parameter and stepsizes
        w.resize(taps, 0);
        mu_f.resize(taps, 0);

        // Exponentially weighted stepsize
        for (int i = 0; i < taps; i++) {
            mu_f[i] = pow(lambda, i) * mu_v;
        }

        // for ring buffer
        for (int i = 1; i < 31; i++) {
            size = 1 << i;
            if ((taps * itvl) <= size)
            break;
        }
        mask = size - 1;

        buf_s.resize(size);
        buf_ns.resize(size);

        srand(1234);
    }

    // for debug
    void
    SemiBlindICA::SBSemiBlindICA::print() const
    {
        printf("tap: %d, mu_s: %f, mu_n: %f\n", taps, mu_s, mu_n);
    }

    //
    void
    SemiBlindICA::SBSemiBlindICA::save(std::ostream &file) const
    {
        float c[2];

        for (int i = 0; i < taps; i++) {
            c[0] = w[i].real();
            c[1] = w[i].imag();
            file.write((const char *)c, sizeof(float) * 2);
        }
    }

    //
    void
    SemiBlindICA::SBSemiBlindICA::load(std::istream &file)
    {
        float c[2];
        for (int i = 0; i < taps; i++) {
            file.read((char *)c, sizeof(float) * 2);
            w[i] = std::complex<float>(c[0], c[1]);
        }
    }

    // filtering
    std::complex<float>
    SemiBlindICA::SBSemiBlindICA::update(const std::complex<float> &x, const std::complex<float> &s)
    {
        std::complex<float> d, sum, c1, n_d;
        int midx;
        float ry, v1, e1, rx, v2, e2;

        //printf("%d/%d, %d/%d\n", ii, id_i, jj, id_j);
        // std::cout << "x=" << x << ", s=" << s << ", w=" << w << std::endl;

        //
        buf_s[idx] = s;

        if (abs(s) < is_zero) {
            return x;
        }

        //
        sum = 0;
        for (int i = 0; i < taps; i++) {
            midx = (idx - i * itvl) & mask;
            sum += w[i] * buf_s[midx];
        }
        d = x - sum;

        //
        n_d = n_var * d;
        buf_ns[idx] = r_var * buf_s[idx];

        // Fixed value
        //srand(1234);

        //
        ry = abs(n_d) + FLOOR_VALUE * rand()/ (double)RAND_MAX;
        v1 = tanh(alpha * ry);
        e1 = ry;
        c1 = v1 * n_d / ry;

        //
        rx = abs(buf_ns[idx]) + FLOOR_VALUE * rand()/(double)RAND_MAX;
        v2 = tanh(alpha * rx);
        e2 = rx;

        // filter update
        for (int i = 0; i < taps; i++) {
            midx = (idx - i * itvl) & mask;
            w[i] += mu_f[i] * (c1 * conj(buf_ns[midx]));
        }

        // variance estimation
        n_var += mu_n * (1 - v1 * e1) * n_var;
        if (n_var < 0) n_var *= -1;
        r_var += mu_s * (1 - v2 * e2) * r_var;
        if (r_var < 0) r_var *= -1;

        idx = (idx + 1) & mask;
        return d;
    }

    SemiBlindICA::SemiBlindICA():
        NodeDef("SemiBlindICA")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("REFERENCE")
            << Terminal::in("CHANNEL").def(1)
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("INTERVAL").def(1)
            << Terminal::in("TAP_LOWERFREQ").def(8)
            << Terminal::in("TAP_UPPERFREQ").def(4)
            << Terminal::in("DECAY").def(0.8)
            << Terminal::in("MU_FILTER").def(0.01)
            << Terminal::in("MU_REFERENCE").def(0.01)
            << Terminal::in("MU_UNKNOWNSIGNAL").def(0.01)
            << Terminal::in("IS_ZERO").def(0.0001)
            << Terminal::in("FILE_FILTER_IN").def("")
            << Terminal::in("FILE_FILTER_OUT").def("")
            << Terminal::in("OUTPUT_FREQ").def(150)
            >> Terminal::out("OUTPUT");
    }

    void
    SemiBlindICA::reset()
    {
        this->nchannel = 0;
        this->length = 0;
        this->itvl = 0;
        this->tap_lower = 0;
        this->tap_upper = 0;
        this->counter_loop = 0;

        this->mu_s = 0;
        this->mu_v = 0;
        this->mu_n = 0;
        this->lambda = 0;
        this->is_zero = 0;

        this->filter_in_str.clear();
        this->filter_out_str.clear();
        this->filter_in_chr = nullptr;
        this->filter_out_chr = nullptr;

        this->counter = 0;
        this->freq = 0;

        this->sbsbica.clear();
    }

    void
    SemiBlindICA::initialize(InputDataBundle &input)
    {
        this->nchannel = input["CHANNEL"].last().as<int>();
        this->length = input["LENGTH"].last().as<int>();
        this->itvl = input["INTERVAL"].last().as<int>();
        this->tap_lower = input["TAP_LOWERFREQ"].last().as<int>();
        this->tap_upper = input["TAP_UPPERFREQ"].last().as<int>();
        this->counter_loop = input["OUTPUT_FREQ"].last().as<int>();

        this->mu_s = input["MU_REFERENCE"].last().as<float>();
        this->mu_v = input["MU_FILTER"].last().as<float>();
        this->mu_n = input["MU_UNKNOWNSIGNAL"].last().as<float>();
        this->lambda = input["DECAY"].last().as<float>();
        this->is_zero = input["IS_ZERO"].last().as<float>();

        this->filter_in_str = input["FILE_FILTER_IN"].last().cast<std::string>();
        this->filter_out_str = input["FILE_FILTER_OUT"].last().cast<std::string>();
        this->filter_in_chr = this->filter_in_str.empty() ? "" : this->filter_in_str.c_str();
        this->filter_out_chr = this->filter_out_str.empty() ? "" : this->filter_out_str.c_str();

        this->counter = 0;
        this->freq = this->length / 2 + 1;

        std::ifstream filter_in_file(this->filter_in_chr, std::ios::binary);
        if(filter_in_file){
            this->read_header(filter_in_file);
        }
    
        this->sbsbica.resize(this->nchannel);
        for (int j = 0; j < this->nchannel; j++) {
            for (int i = 0; i < this->freq; i++) {
                // linear interpolation of tap length
                int taps = (int)((double)this->tap_lower
                        + (double)(this->tap_upper - this->tap_lower)
                        * (double)i / (double)this->freq);
                // add element
                this->sbsbica[j].emplace_back(taps, itvl, lambda, mu_s, mu_v, mu_n, is_zero);
                if(filter_in_file){
                    this->sbsbica[j].back().load(filter_in_file);
                }
            }
        }
    }

    void
    SemiBlindICA::read_header(std::istream &file)
    {
        int value;
        file.read((char *)&value, sizeof(int));
        if(value != this->nchannel){
            throw RuntimeError("Channel count mismatch");
        }
        file.read((char *)&value, sizeof(int));
        if(value != this->freq){
            throw RuntimeError("Freq mismatch");
        }
        file.read((char *)&value, sizeof(int));
        if(value != this->tap_upper){
            throw RuntimeError("Tap upper mismatch");
        }
        file.read((char *)&value, sizeof(int));
        if(value != this->tap_lower){
            throw RuntimeError("Tap lower mismatch");
        }
    }

    void
    SemiBlindICA::write_header(std::ostream &file) const
    {
        file.write((const char *)&this->nchannel, sizeof(int));
        file.write((const char *)&this->freq, sizeof(int));
        file.write((const char *)&this->tap_upper, sizeof(int));
        file.write((const char *)&this->tap_lower, sizeof(int));
    }

    bool
    SemiBlindICA::operator()(size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->initialize(input);
        }

        const auto &input1 = input["INPUT"].last().cast<Matrix<std::complex<float>>>();
        const auto &input2 = input["REFERENCE"].last().cast<Matrix<std::complex<float>>>();

        // {
        //     std::stringstream ss;
        //     ss << "SemiBlindICA: count=" << count << std::endl;
        //     ss << "input1: size=(" << input1.rows() << ", " << input1.cols() << "), values=" << input1 << std::endl;
        //     ss << "input2: size=(" << input2.rows() << ", " << input2.cols() << "), values=" << input2 << std::endl;
        //     std::cout << ss.str();
        // }

        const int row = input1.rows();
        const int col = input1.cols();

        auto outData = data::create<Matrix<std::complex<float>>>(row, col);
        auto &out = outData.cast<Matrix<std::complex<float>>>();

        // update filter
        // std::cout << "ch=" << this->nchannel << ", freq=" << this->freq << std::endl;
        for (int ch = 0; ch < this->nchannel; ch++) {
            for (int w = 0; w < this->freq; w++) {
                out(ch, w) = sbsbica[ch][w].update(input1(ch, w), input2(0, w));
            }
        }

        // save filter
        this->counter++;
        if (this->filter_out_chr != NULL && this->counter >= this->counter_loop) {
            std::ofstream filter_out_file(filter_out_chr, std::ios::binary);
            this->write_header(filter_out_file);
            for (int ch = 0; ch < this->nchannel; ch++) {
                for (int w = 0; w < this->freq; w++) {
                    sbsbica[ch][w].save(filter_out_file);
                }
            }
            this->counter = 0;
        }

        // {
        //     for(size_t i = 0; i < out.rows(); i++){
        //         for(size_t j = 0; j < out.cols(); j++){
        //             std::stringstream ss;
        //             ss << "SemiBlindICA: output: " << count << ", " << i << ", " << j << ", " << out(i, j) << std::endl;
        //             std::cout << ss.str();
        //         }
        //     }
        // }
        output["OUTPUT"].push(outData);
        return true;
    }
}
}
