// -*-mode:c++; c-style:k&r; c-basic-offset:4;-*-
//
// Copyright 2011-2021, Julian Catchen <jcatchen@illinois.edu>
//
// This file is part of Stacks.
//
// Stacks is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Stacks is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Stacks.  If not, see <http://www.gnu.org/licenses/>.
//

#ifndef __POPSUM_H__
#define __POPSUM_H__

#include <cstring>
#include <string>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <utility>
#include <cstdint>
#include <cmath>

#include "stacks.h"
#include "locus.h"
#include "PopMap.h"
#include "MetaPopInfo.h"
#include "Hwp.h"

extern bool      calc_hwp;
extern double    minor_allele_freq;
extern corr_type fst_correction;
extern double    p_value_cutoff;

const uint16_t PopStatSize = 6;

class PopStat {
public:
    int      loc_id;
    int      bp;
    bool     fixed;
    uint16_t alleles;    // Number of alleles sampled at this location.

    PopStat() {
        this->loc_id   = 0;
        this->bp       = 0;
        this->fixed    = false;
        this->alleles  = 0;
    }
    int reset() {
        this->loc_id   = 0;
        this->bp       = 0;
        this->fixed    = false;
        this->alleles  = 0;

        return 0;
    }
    void optimize() {
	cerr << "  popstat.loc_id:  " << sizeof(this->loc_id) << "\n"
	     << "  popstat.bp:      " << sizeof(this->bp) << "\n"
	     << "  popstat.fixed:   " << sizeof(this->fixed) << "\n"
	     << "  popstat.alleles: " << sizeof(this->alleles) << "\n";
    }
    virtual ~PopStat() {}
};

class HapStat: public PopStat {
public:
    static const uint8_t _hapstat_size = 6;
    // stat[0]: Phi_st
    // stat[1]: Fst'
    // stat[2]: D_est
    // stat[3]: Dxy
    // stat[4]: Phi_ct
    // stat[5]: Phi_sc

private:
    double _stat[_hapstat_size];
    double _smoothed[_hapstat_size];
    double _bs[_hapstat_size];

public:
    uint16_t  pop_1;
    uint16_t  pop_2;
    uint16_t  popcnt;

    HapStat(): PopStat() {
        popcnt = uint16_t(-1);

	for (uint i = 0; i < this->_hapstat_size; i++) {
            this->_stat[i]     =  0.0;
            this->_smoothed[i] = -1.0;
            this->_bs[i]       = -1.0;
        }
    }
    int reset() {
        PopStat::reset();
	for (uint i = 0; i < this->_hapstat_size; i++) {
            this->_stat[i]     =  0.0;
            this->_smoothed[i] = -1.0;
            this->_bs[i]       = -1.0;
        }
	return 0;
    }

    static uint16_t size() { return _hapstat_size; };
    double stat(uint16_t index) const { assert(index < this->_hapstat_size); return this->_stat[index]; }
    void   stat(uint16_t index, double value) { assert(index < this->_hapstat_size); this->_stat[index] = value; }
    double smoothed(uint16_t index) const { assert(index < this->_hapstat_size); return this->_smoothed[index]; }
    void   smoothed(uint16_t index, double value) { assert(index < this->_hapstat_size); this->_smoothed[index] = value; }
    double bs(uint16_t index) const { assert(index < this->_hapstat_size); return this->_bs[index]; }
    void   bs(uint16_t index, double value) { assert(index < this->_hapstat_size); this->_bs[index] = value; }
};

class LocStat: public PopStat {
public:
    static const uint8_t _locstat_size = 4;
    // stat[0]: gene diversity
    // stat[1]: haplotype diversity (Pi)
    // stat[2]: Hardy-Weinberg proportions p-value
    // stat[3]: Hardy-Weinberg proportions p-value standard error.
    
private:
    double _stat[_locstat_size];
    double _smoothed[_locstat_size];
    double _bs[_locstat_size];

public:
    uint16_t hap_cnt; // Number of unique haplotypes at this locus.
    string   hap_str; // Human-readable string of haplotype counts.

    LocStat(): PopStat() {
        this->hap_cnt = 0;

	for (uint i = 0; i < this->_locstat_size; i++) {
            this->_stat[i]     =  0.0;
            this->_smoothed[i] = -1.0;
            this->_bs[i]       = -1.0;
        }
    }
    ~LocStat() {};
    int reset() {
        PopStat::reset();
	for (uint i = 0; i < this->_locstat_size; i++) {
            this->_stat[i]     =  0.0;
            this->_smoothed[i] = -1.0;
            this->_bs[i]       = -1.0;
        }
	return 0;
    }

    static uint16_t size() { return _locstat_size; };
    double stat(uint16_t index) const { assert(index < this->_locstat_size); return this->_stat[index]; }
    void   stat(uint16_t index, double value) { assert(index < this->_locstat_size); this->_stat[index] = value; }
    double smoothed(uint16_t index) const { assert(index < this->_locstat_size); return this->_smoothed[index]; }
    void   smoothed(uint16_t index, double value) { assert(index < this->_locstat_size); this->_smoothed[index] = value; }
    double bs(uint16_t index) const { assert(index < this->_locstat_size); return this->_bs[index]; }
    void   bs(uint16_t index, double value) { assert(index < this->_locstat_size); this->_bs[index] = value; }
};

class PopPair: public PopStat {
public:
    static const uint8_t _poppair_size = 1;
    // stat[0]: AMOVA Fst

private:
    double   _stat[_poppair_size];
    double   _smoothed[_poppair_size];
    double   _bs[_poppair_size];

public:
    uint16_t pop_1;
    uint16_t pop_2;
    uint16_t col;
    double   fet_p;      // Fisher's Exact Test p-value.
    float    fet_or;     // Fisher's exact test odds ratio.
    float    or_se;      // Fisher's exact test odds ratio standard error.
    float    lod;        // base 10 logarithm of odds score.
    float    ci_low;     // Fisher's exact test lower confidence interval.
    float    ci_high;    // Fisher's exact test higher confidence interval.
    
    PopPair(): PopStat() {
        pop_1     = 0;
        pop_2     = 0;
        col       = 0;
        fet_p     = 0.0;
        fet_or    = 0.0;
        or_se     = 0.0;
        lod       = 0.0;
        ci_low    = 0.0;
        ci_high   = 0.0;

	for (uint i = 0; i < this->_poppair_size; i++) {
            this->_stat[i]     =  0.0;
            this->_smoothed[i] = -1.0;
            this->_bs[i]       = -1.0;
        }
    }
    void optimize() {
	cerr << "Size of PopPair:    " << sizeof(*this) << "\n";
	PopStat::optimize();
	cerr << "  poppair.pop_1:    " << sizeof(this->pop_1) << "\n"
	     << "  poppair.pop_2:    " << sizeof(this->pop_2) << "\n"
	     << "  poppair.col:      " << sizeof(this->col) << "\n"
	     << "  poppair.fet_p:    " << sizeof(this->fet_p) << "\n"
	     << "  poppair.fet_or:   " << sizeof(this->fet_or) << "\n"
	     << "  poppair.or_se:    " << sizeof(this->or_se) << "\n"
             << "  poppair.lod:      " << sizeof(this->lod) << "\n"
             << "  poppair.ci_low:   " << sizeof(this->ci_low) << "\n"
             << "  poppair.ci_high:  " << sizeof(this->ci_high) << "\n"
	     << "  poppair.stat:     " << sizeof(this->_stat) << "\n"
	     << "  poppair.smoothed: " << sizeof(this->_smoothed) << "\n"
	     << "  poppair.bs:       " << sizeof(this->_bs) << "\n";
    }

    static uint16_t size() { return _poppair_size; };
    double stat(uint16_t index) const { assert(index < this->_poppair_size); return this->_stat[index]; }
    void   stat(uint16_t index, double value) { assert(index < this->_poppair_size); this->_stat[index] = value; }
    double smoothed(uint16_t index) const { assert(index < this->_poppair_size); return this->_smoothed[index]; }
    void   smoothed(uint16_t index, double value) { assert(index < this->_poppair_size); this->_smoothed[index] = value; }
    double bs(uint16_t index) const { assert(index < this->_poppair_size); return this->_bs[index]; }
    void   bs(uint16_t index, double value) { assert(index < this->_poppair_size); this->_bs[index] = value; }

    // AMOVA Fst method, from Weir, Genetic Data Analysis II.
    double   amova_fst() const       { return this->_stat[0]; }
    void     amova_fst(double value) { this->_stat[0] = value; }
};

class SumStat: public PopStat {
    static const uint8_t _sumstat_size = 3;
    // stat[0]: pi
    // stat[1]: fis
    // stat[2]: HWE p-value

    struct HetNuc {
        char   _q_nuc;
        double _p;
        float  _obs_het;
        float  _obs_hom;
        float  _exp_het;
        float  _exp_hom;
        double _stat[_sumstat_size];
        double _smoothed[_sumstat_size];
        double _bs[_sumstat_size];
    };
    bool     _filtered_site;
    uint16_t _num_indv;
    char     _p_nuc;
    HetNuc  *_het_nuc;

public:
    SumStat(): PopStat() {
        _het_nuc       = NULL;
        _num_indv      = 0;
        _p_nuc         = 0;
        _filtered_site = false;
    }
    ~SumStat() {
        if (this->_het_nuc != NULL)
            delete this->_het_nuc;
    }
    int reset() {
        PopStat::reset();

        if (_het_nuc != NULL) {
            delete _het_nuc;
            _het_nuc = NULL;
        }
        
        _num_indv      = 0;
        _p_nuc         = 0;
        _filtered_site = false;

        return 0;
    }
    void optimize() {
	cerr << "Size of SumStat:    " << sizeof(*this) << "\n";
	PopStat::optimize();
	cerr << "  sumstat.filtered_site: " << sizeof(this->_filtered_site) << "\n"
	     << "  sumstat.num_indv:  " << sizeof(this->_num_indv) << "\n"
             << "  sumstat._het_nuc:  " << sizeof(this->_het_nuc) << "\n"
             << "  *sumstat._het_nuc: " << (this->_het_nuc ? sizeof(*this->_het_nuc) : 0) << "\n"
             << "  sumstat.p_nuc:     " << sizeof(this->_p_nuc)    << "\n";
    }

    bool variable_site() { return this->_het_nuc == NULL ? false : true; }
    void variable_site(bool v) {
        if (v) {
            this->_het_nuc = new HetNuc;

            for (uint i = 0; i < this->_sumstat_size - 1; i++) {
                this->_het_nuc->_stat[i]     =  0.0;
                this->_het_nuc->_smoothed[i] = -1.0;
                this->_het_nuc->_bs[i]       = -1.0;
            }
            // The default value for the HWE p-value should be 1.0, not 0.
            this->_het_nuc->_stat[this->_sumstat_size - 1] = 1.0; 
        } else {
            if (this->_het_nuc != NULL) {
                delete this->_het_nuc;
                this->_het_nuc = NULL;
            }
        }
    }
    uint16_t num_indv()             { return this->_num_indv; }
    void     num_indv(uint16_t n)   { this->_num_indv = n; }
    double   p()                    { if (this->_het_nuc != NULL) return this->_het_nuc->_p; else return 1.0; }
    void     p(double p)            { if (this->_het_nuc != NULL) this->_het_nuc->_p = p; }
    char     p_nuc()                { return this->_p_nuc; }
    void     p_nuc(char p_nuc)      { this->_p_nuc = p_nuc; }
    char     q_nuc()                { if (this->_het_nuc != NULL) return this->_het_nuc->_q_nuc; else return 0; }
    void     q_nuc(char q_nuc)      { if (this->_het_nuc != NULL) this->_het_nuc->_q_nuc = q_nuc; }
    float    obs_het()              { if (this->_het_nuc != NULL) return this->_het_nuc->_obs_het; else return 0.0; }
    void     obs_het(float obs_het) { if (this->_het_nuc != NULL) this->_het_nuc->_obs_het = obs_het; }
    float    obs_hom()              { if (this->_het_nuc != NULL) return this->_het_nuc->_obs_hom; else return 1.0; }
    void     obs_hom(float obs_hom) { if (this->_het_nuc != NULL) this->_het_nuc->_obs_hom = obs_hom; }
    float    exp_het()              { if (this->_het_nuc != NULL) return this->_het_nuc->_exp_het; else return 0.0; }
    void     exp_het(float exp_het) { if (this->_het_nuc != NULL) this->_het_nuc->_exp_het = exp_het; }
    float    exp_hom()              { if (this->_het_nuc != NULL) return this->_het_nuc->_exp_hom; else return 1.0; }
    void     exp_hom(float exp_hom) { if (this->_het_nuc != NULL) this->_het_nuc->_exp_hom = exp_hom; }
    bool     filtered_site()        { return this->_filtered_site; }
    void     filtered_site(bool v)  { this->_filtered_site = v; }

    static uint16_t size() { return _sumstat_size; }
    
    double stat(uint16_t index) const {
        assert(index < this->_sumstat_size);
        if (this->_het_nuc != NULL) {
            return this->_het_nuc->_stat[index];
        } else {
            switch(index) {
            case 0: // Pi
                return 0.0;
            case 1: // Fis
                return FSNC;
            case 2: // HWE p-value
                return 1.0;
            default:
                DOES_NOT_HAPPEN;
                return 0.0;
            }
        }
    }
    void   stat(uint16_t index, double value)     { if (this->_het_nuc != NULL) { assert(index < this->_sumstat_size); this->_het_nuc->_stat[index] = value; } }
    double smoothed(uint16_t index) const         { if (this->_het_nuc != NULL) { assert(index < this->_sumstat_size); return this->_het_nuc->_smoothed[index]; } else return 0.0; }
    void   smoothed(uint16_t index, double value) { if (this->_het_nuc != NULL) { assert(index < this->_sumstat_size); this->_het_nuc->_smoothed[index] = value; } }
    double bs(uint16_t index) const               { if (this->_het_nuc != NULL) { assert(index < this->_sumstat_size); return this->_het_nuc->_bs[index]; } else return 0.0; }
    void   bs(uint16_t index, double value)       { if (this->_het_nuc != NULL) { assert(index < this->_sumstat_size); this->_het_nuc->_bs[index] = value; } }

    double pi() const          { if (this->_het_nuc != NULL) { return this->_het_nuc->_stat[0]; } else return 0.0; }
    void   pi(double value)    { if (this->_het_nuc != NULL) { this->_het_nuc->_stat[0] = value; } }
    double fis() const         { if (this->_het_nuc != NULL) { return this->_het_nuc->_stat[1]; } else return FSNC; }
    void   fis(double value)   { if (this->_het_nuc != NULL) { this->_het_nuc->_stat[1] = value; } }
    double hwe_p() const       { if (this->_het_nuc != NULL) { return this->_het_nuc->_stat[2]; } else return 1.0; }
    void   hwe_p(double value) { if (this->_het_nuc != NULL) { this->_het_nuc->_stat[2] = value; } }
};

class LocSum {
public:
    //
    // Array containing summary statistics for each nucleotide position at this locus.
    //
    SumStat *nucs;
    LocSum(int len) {
        this->nucs = new SumStat[len];
    }
    ~LocSum() {
        delete [] this->nucs;
    }
    void optimize(const size_t loc_len) {
        cerr << "Size of LocSum: " << sizeof(*this) << "\n";
        uint s = 0;
        for (uint i = 0; i < loc_len; i++)
            s += sizeof(this->nucs[i]);
        cerr << "  locsum.nucs:  " << s << "\n";
    }
    size_t memusage(const size_t loc_len) const {
        size_t s = 0;
        s += sizeof(*this);
        for (uint i = 0; i < loc_len; i++)
            s += sizeof(this->nucs[i]);
        return s;
    }
};

class NucTally {
public:
    int      loc_id;
    int      bp;
    uint16_t col;
    uint16_t num_indv;
    uint16_t pop_cnt;
    uint16_t allele_cnt;
    char     p_allele;
    char     q_allele;
    float    p_freq;
    float    obs_het;
    bool     fixed;
    int16_t  priv_allele;

    NucTally() {
        loc_id      = 0;
        bp          = 0;
        col         = 0;
        num_indv    = 0;
        pop_cnt     = 0;
        allele_cnt  = 0;
        p_allele    = 0;
        q_allele    = 0;
        p_freq      = 0.0;
        obs_het     = 0.0;
        priv_allele = -1;
        fixed       = true;
    }
    int reset() {
        loc_id      = 0;
        bp          = 0;
        col         = 0;
        num_indv    = 0;
        pop_cnt     = 0;
        allele_cnt  = 0;
        p_allele    = 0;
        q_allele    = 0;
        p_freq      = 0.0;
        obs_het     = 0.0;
        priv_allele = -1;
        fixed       = true;
        return 0;
    }
    void  optimize() {
	cerr << "Size of NucTally:    " << sizeof(*this)          << "\n"
	     << "  nuctally.loc_id:   " << sizeof(this->loc_id)   << "\n"
	     << "  nuctally.bp:       " << sizeof(this->bp)       << "\n"
	     << "  nuctally.col:      " << sizeof(this->col)      << "\n"
	     << "  nuctally.num_indv: " << sizeof(this->num_indv) << "\n"
	     << "  nuctally.pop_cnt:  " << sizeof(this->pop_cnt)  << "\n"
	     << "  nuctally.allele_cnt:  " << sizeof(this->allele_cnt) << "\n"
	     << "  nuctally.p_allele: " << sizeof(this->p_allele) << "\n"
	     << "  nuctally.q_allele: " << sizeof(this->q_allele) << "\n"
	     << "  nuctally.p_freq:   " << sizeof(this->p_freq)   << "\n"
	     << "  nuctally.obs_het:  " << sizeof(this->obs_het)  << "\n"
	     << "  nuctally.fixed:    " << sizeof(this->fixed)    << "\n"
	     << "  nuctally.priv_allele: " << sizeof(this->priv_allele) << "\n";
    }
};

class LocTally {
public:
    NucTally *nucs;

    LocTally(int len)  {
        this->nucs = new NucTally[len];
    }
    ~LocTally() {
        delete [] this->nucs;
    }
};

//
// per-Locus Population-level summary
//
class LocPopSum {
    size_t    _pop_cnt;
    LocSum  **_per_pop;
    LocTally *_meta_pop;
    LocStat **_hapstats_per_pop;

public:
    LocPopSum(size_t cloc_len, const MetaPopInfo& mpopi);
    ~LocPopSum();

    int             sum_pops(const CSLocus *, Datum const*const*, const MetaPopInfo&, bool, ostream &);
    int             tally_metapop(const CSLocus *);
    int             calc_hapstats(const CSLocus *, const Datum **, const MetaPopInfo&);
    const LocSum  **all_pops()                         { return (const LocSum **) this->_per_pop; }
    const LocSum   *per_pop(size_t pop_index)          const { return this->_per_pop[pop_index]; }
    const LocTally *meta_pop()                         const { return this->_meta_pop; }
    LocStat        *hapstats_per_pop(size_t pop_index) const { return this->_hapstats_per_pop[pop_index]; }
    size_t          pop_cnt()                          const { return this->_pop_cnt; }
    static double   pi(double, double, double);
    static double   binomial_coeff(double, double);
    double          hwe(double, double, double, double, double, double);

private:
    int      tally_heterozygous_pos(const CSLocus *, Datum const*const*, LocSum *, int, int, uint, uint);
    int      tally_fixed_pos(const CSLocus *, Datum const*const*, LocSum *, int, uint, uint);
    int      tally_ref_alleles(int, uint16_t &, char &, char &, uint16_t &, uint16_t &);
    int      tally_observed_haplotypes(const vector<char *> &, int);
    LocStat *haplotype_diversity(int, int, Datum const*const*);
    double   log_hwp_pr(double, double, double, double, double, double);
};

struct LocBin {
    size_t     sample_cnt;
    CSLocus   *cloc;
    Datum    **d;
    LocPopSum *s;

    LocBin(size_t cnt): sample_cnt(cnt), cloc(NULL), d(NULL), s(NULL) {}
    ~LocBin() {
        if (this->cloc != NULL) delete cloc;
        if (this->d    != NULL) {
            for (uint i = 0; i < sample_cnt; i++)
                if (d[i] != NULL) delete d[i];
            delete [] d;
        }
        if (this->s    != NULL) delete s;
    }
};

//
// per-Locus class for calculating divergence values such as Fst.
//
class LocusDivergence {
    const MetaPopInfo         *_mpopi;
    vector<vector<PopPair **>> _snps;
    vector<vector<HapStat *>>  _haplotypes;
    vector<HapStat *>          _metapop_haplotypes;

    vector<double> _mean_fst;
    vector<double> _mean_fst_cnt;
    vector<double> _mean_phist;
    vector<double> _mean_phist_cnt;
    vector<double> _mean_fstp;
    vector<double> _mean_fstp_cnt;
    vector<double> _mean_dxy;
    vector<double> _mean_dxy_cnt;

public:
    LocusDivergence(const MetaPopInfo *mpopi);
    int clear(const vector<LocBin *> &loci);
    int snp_divergence(const vector<LocBin *> &loci);
    int haplotype_divergence_pairwise(const vector<LocBin *> &loci);
    int haplotype_divergence(const vector<LocBin *> &loci);

    int write_summary(string);

    vector<vector<PopPair **>>& snp_values()       { return this->_snps; }
    vector<vector<HapStat *>>&  haplotype_values() { return this->_haplotypes; }
    vector<HapStat *>&  metapop_haplotype_values() { return this->_metapop_haplotypes; }

private:
    //
    // SNP-level F statistics.
    //
    PopPair *Fst(const CSLocus *, const LocPopSum *, int, int, int);
    int      fishers_exact_test(PopPair *, double, double, double, double);

    //
    // Haplotype-level F statistics
    //
    double   haplotype_d_est(const Datum **, const LocSum **, vector<int> &);
    HapStat *haplotype_amova(const Datum **, const LocSum **, vector<int> &);
    double   amova_ssd_total(vector<string> &, map<string, int> &, double **);
    double   amova_ssd_wp(vector<int> &, map<int, vector<int>> &, map<string, int> &, map<int, vector<string>> &, double **);
    double   amova_ssd_ap_wg(vector<int> &, map<int, vector<int>> &, map<string, int> &, map<int, vector<string>> &, double **, double **);
    double   amova_ssd_ag(vector<int> &, map<int, vector<int>> &, map<string, int> &, map<int, vector<string>> &, double **, double);

    //
    // Haplotype-level Dxy
    //
    double   haplotype_dxy(const Datum **, size_t,  vector<int> &);

    bool     fixed_locus(const Datum **, vector<int> &);
    int      nuc_substitution_identity(map<string, int> &, double **);
    int      nuc_substitution_identity_max(map<string, int> &, double **);
};

//
// Utility functions for summary/haplotype statistics.
//
bool     uncalled_haplotype(const char *);
double   count_haplotypes_at_locus(int, int, Datum const*const*, map<string, double> &);
int      nuc_substitution_dist(map<string, int> &, double **);


#endif // __POPSUM_H__
