"""
Modified by Matthias Rüegg (matthias.ruegg@unil.ch) on August 13 2019
Based on a program created by created by github.com/blmoistawinde (https://github.com/blmoistawinde/facetnet-python)
Copyright 2019 __UNIL__. All rights reserved.

This file is part of facetnet-python-unil.

    facetnet-python-unil 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.

    facetnet-python-unil 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 facetnet-python-unil. If not, see <https://www.gnu.org/licenses/>.
"""

import argparse
import sys
import numpy as np
import facetnet as fn
import networkx as nx
import community as cy
import matplotlib.pyplot as plt

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("edgelist", help="full path to edge-list file", type=str)
    parser.add_argument("res_folder", help="path to results folder, where"
                                           "xcap.dat, lcap.dat, idmap.dat, idmap_inv.dat, soft_comm.csv"
                                           "will be saved", type=str)
    parser.add_argument("--show", help="show convergence plot of gamma", type=str)
    args = parser.parse_args()

    print("Analyzing network: {}".format(args.edgelist))

    if args.show is not None:
        is_show = True
    else:
        is_show = False

    max_comm = 7
    Q = []
    Qs = []

    for m in range(2, max_comm):
        print("# clusters: {}".format(m))
        idmap, idmap_inv, xc, lc, qc_s, soft_comm_df, comm_net_df, evol_net_df = fn.facetnet_step(
            args.edgelist, 1.0, m, show_plot=is_show, weighted=False)
        Qs.append(qc_s)

        graph = nx.read_edgelist(args.edgelist, nodetype=int)

        soft_comm_df.to_csv("{}/soft_comm_m{}.csv".format(args.res_folder, m), index=False)
        comm_net_df.to_csv("{}/comm_net_m{}.csv".format(args.res_folder, m), index=False)

        hard_comm = soft_comm_df.drop('id', axis=1).idxmax(axis=1)
        cluster_labels = list(soft_comm_df.drop('id', axis=1).columns)
        partition = dict(zip(soft_comm_df['id'], hard_comm))
        for node in partition:
            partition[node] = cluster_labels.index(partition[node])

        Q.append(cy.modularity(partition, graph))

    plt.plot(np.arange(2, max_comm, 1), Q, 'ro-')
    plt.plot(np.arange(2, max_comm, 1), Qs, 'bs-')
    plt.legend(["Q", "Qs"])
    plt.xticks(np.arange(2, max_comm, 1))
    plt.ylim([0.0, 0.6])
    plt.xlabel("Community number")
    plt.show()
    print("Results written to: {}".format(args.res_folder))