"""Module that uses INSPIRE-HEP to get author
(number of papers, h index, citations) and paper statistics (citations).

Uses matplotlib to do plots.

This file is part of the physbiblio package.
"""
import time
import os
import traceback
import os.path as osp
import json
import requests
import dateutil
import datetime
import matplotlib

matplotlib.use("Qt5Agg")
os.environ["QT_API"] = "pyside2"
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.backends.backend_pdf import PdfPages

try:
    from physbiblio.errors import pBLogger
    from physbiblio.config import pbConfig
except ImportError:
    print("Could not find physbiblio and its modules!")
    print(traceback.format_exc())
    raise


class InspireStatsLoader:
    """Class that contains the methods
    to collect information from INSPIRE-HEP
    """

    urlBase = pbConfig.inspireSearchBase
    timeout = float(pbConfig.params["timeoutWebSearch"])
    authorStatsOpts = "&of=recjson&ot=recid,creation_date&so=a&rg=250"
    paperStatsOpts = "&of=recjson&ot=recid,creation_date&so=a&rg=250"
    skipPageOpt = "&jrec="
    maxPerPage = 250
    authorPlotInfo = {}
    paperPlotInfo = {}

    def __init__(self):
        """The class constructor,
        defines some constants and search options
        """
        self.urlBase = pbConfig.inspireSearchBase
        self.timeout = float(pbConfig.params["timeoutWebSearch"])
        self.authorPlotInfo = None
        self.paperPlotInfo = None
        self.allInfoA = {}
        self.authorPapersList = [[], []]
        self.allCitations = []
        self.runningAuthorStats = True
        self.runningPaperStats = True
        self.allInfoP = {}
        self.citingPapersList = [[], []]

    def changeBackend(self, wantBackend):
        """Changes the matplotlib backend currently in use.

        Parameters:
            wantBackend: a string that defines the wanted backend
        """
        if wantBackend != matplotlib.get_backend():
            matplotlib.use(wantBackend, warn=False, force=True)
            from matplotlib import pyplot as plt

            pBLogger.info("Changed backend to %s" % matplotlib.get_backend())

    def JsonFromUrl(self, url):
        """Function that downloads the url content
        and returns a Json object

        Parameters:
            url: string containing the url to be opened

        Output:
            the json object generated from the url content
        """

        def getSeries(url):
            response = requests.get(url, timeout=self.timeout)
            try:
                return json.loads(response.content.decode("utf-8"))
            except ValueError:
                pBLogger.warning("Empty response!")
                return []
            except:
                pBLogger.exception("Cannot read the page content!")
                return []

        ser = 0
        complete = []
        while True:
            temp = getSeries(
                url + self.skipPageOpt + "%d" % (ser * self.maxPerPage + 1)
            )
            if len(temp) < self.maxPerPage:
                return complete + temp
            else:
                complete += temp
                ser += 1

    def authorStats(self, authorName, plot=False, reset=True, pbMax=None, pbVal=None):
        """Function that gets the data and
        constructs the statistics for a given author.

        Parameters:
            authorName: the author name as identified into INSPIRE-HEP,
                or a list of author names
                (it calls itself recursively for all the list elements)
            plot (boolean, default False): True to call self.plotStats
            reset (boolean, default False): True to delete
                all previous existing data
                (used as False when processing a list of authors)
            pbMax (callable, optional): a function to set the maximum
                of a progress bar in the GUI, if possible
            pbVal (callable, optional): a function to set the value
                of a progress bar in the GUI, if possible

        Output:
            a dictionary containing all the statistic information.
            For a single author, the structure is the following:
            {
                "name": the author name,
                "aI": The complete information,
                    including the dictionaries with the single papers info
                    (see self.paperStats), the citations
                    and the corresponding dates,
                "paLi": a list of [id, date] of the papers
                    associated with the author,
                "allLi": the complete list of [date, total citations]
                    with all the citations to the papers,
                "meanLi": the complete list of
                    [date, total citations/number of papers]
                    computed at each point from "allLi" content,
                "h": the h-index,
                "figs" (only if `plot` is True): contains the figures.
                    See self.plotStats
            }
        """
        if reset:
            self.allInfoA = {}
            self.authorPapersList = [[], []]
            self.allCitations = []
        if isinstance(authorName, list):
            try:
                pbMax(len(authorName))
            except TypeError:
                pass
            for ia, a in enumerate(authorName):
                try:
                    pbVal(ia + 1)
                except TypeError:
                    pass
                self.authorStats(a, reset=False)
            self.authorPlotInfo["name"] = authorName
            return self.authorPlotInfo
        pBLogger.info("Stats for author '%s'" % authorName)
        url = (
            pbConfig.inspireSearchBase
            + "?p=author:"
            + authorName
            + self.authorStatsOpts
        )
        data = self.JsonFromUrl(url)
        recid_authorPapers = ["%d" % a["recid"] for a in data]
        tot = len(recid_authorPapers)
        pBLogger.info(
            "AuthorStats will process %d " % tot + "total papers to retrieve citations"
        )
        self.runningAuthorStats = True
        for i, p in enumerate(recid_authorPapers):
            if not self.runningAuthorStats:
                pBLogger.info(
                    "Received 'stop' signal. "
                    + "Interrupting download of information. "
                    + "Processing and exiting..."
                )
                break
            if tot > 20:
                time.sleep(1)
            if p in self.allInfoA.keys():
                continue
            self.allInfoA[p] = {}
            self.allInfoA[p]["date"] = dateutil.parser.parse(data[i]["creation_date"])
            self.authorPapersList[0].append(self.allInfoA[p]["date"])
            pBLogger.info(
                "%5d / %d (%5.2f%%) - looking for paper: '%s'"
                % (i + 1, tot, 100.0 * (i + 1) / tot, p)
            )
            paperInfo = self.paperStats(
                p, verbose=0, paperDate=self.allInfoA[p]["date"]
            )
            self.allInfoA[p]["infoDict"] = paperInfo["aI"]
            self.allInfoA[p]["citingPapersList"] = paperInfo["citList"]
            for c, v in self.allInfoA[p]["infoDict"].items():
                self.allCitations.append(v["date"])
            pBLogger.info("")
        self.authorPapersList[1] = []
        for i, p in enumerate(sorted(self.authorPapersList[0])):
            self.authorPapersList[0][i] = p
            self.authorPapersList[1].append(i + 1)
        pBLogger.info("Saving citation counts...")
        allCitList = [[], []]
        meanCitList = [[], []]
        currPaper = 0
        for i, d in enumerate(sorted(self.allCitations)):
            if (
                currPaper < len(self.authorPapersList[0]) - 1
                and d >= self.authorPapersList[0][currPaper + 1]
            ):
                currPaper += 1
            allCitList[0].append(d)
            allCitList[1].append(i + 1)
            meanCitList[0].append(d)
            meanCitList[1].append((i + 1.0) / self.authorPapersList[1][currPaper])
        hind = 0
        citations = [
            len(self.allInfoA[k]["citingPapersList"][0]) - 2
            for k in self.allInfoA.keys()
        ]
        for h in range(len(citations)):
            if len([a for a in citations if a >= h]) >= h:
                hind = h
        self.authorPlotInfo = {
            "name": authorName,
            "aI": self.allInfoA,
            "paLi": self.authorPapersList,
            "allLi": allCitList,
            "meanLi": meanCitList,
            "h": hind,
        }
        if plot:
            self.authorPlotInfo["figs"] = self.plotStats(author=True)
        pBLogger.info("Stats for author '%s' completed!" % authorName)
        return self.authorPlotInfo

    def paperStats(
        self,
        paperID,
        plot=False,
        verbose=1,
        paperDate=None,
        reset=True,
        pbMax=None,
        pbVal=None,
    ):
        """Function that gets the data and
        constructs the statistics for a given paper.

        Parameters:
            paperID (string): the INSPIRE-HEP id of the paper (a number)
            plot (boolean): whether or not the citations
                should be plotted (default False)
            verbose (int, default 1): increase the verbosity level
            paperDate (datetime, optional): the date of at which
                the paper was published
            reset (boolean, default False): True to delete
                all previous existing data
                (used as False when processing a list of IDs)
            pbMax (callable, optional): a function to set the maximum
                of a progress bar in the GUI, if possible
            pbVal (callable, optional): a function to set the value
                of a progress bar in the GUI, if possible

        Output:
            a dictionary containing all the desired information.
            The structure is the following:
            {
                "id": the paper ID,
                "aI": the list of creation date for all the papers,
                    in INSPIRE-HEP order,
                "citList": the ordered list of citing papers,
                "fig" (only if `plot` is True): contains the figure.
                    See self.plotStats
            }
        """
        if reset:
            self.allInfoP = {}
            self.citingPapersList = [[], []]
        if isinstance(paperID, list):
            self.runningPaperStats = True
            try:
                pbMax(len(paperID))
            except TypeError:
                pass
            for ia, a in enumerate(paperID):
                try:
                    pbVal(ia + 1)
                except TypeError:
                    pass
                if self.runningPaperStats:
                    self.paperStats(a, reset=False)
            self.paperPlotInfo["id"] = paperID
            return self.paperPlotInfo
        if verbose > 0:
            pBLogger.info("Stats for paper '%s'" % paperID)
        url = (
            pbConfig.inspireSearchBase
            + "?p=refersto:recid:"
            + paperID
            + self.paperStatsOpts
        )
        data = self.JsonFromUrl(url)
        recid_citingPapers = [a["recid"] for a in data]
        if paperDate is not None:
            self.citingPapersList[0].append(paperDate)
        for i, p in enumerate(recid_citingPapers):
            self.allInfoP[p] = {}
            self.allInfoP[p]["date"] = dateutil.parser.parse(data[i]["creation_date"])
            self.citingPapersList[0].append(self.allInfoP[p]["date"])
        for i, p in enumerate(sorted(self.citingPapersList[0])):
            self.citingPapersList[0][i] = p
            self.citingPapersList[1].append(i + 1)
        self.citingPapersList[0].append(
            datetime.datetime.fromordinal(datetime.date.today().toordinal())
        )
        try:
            self.citingPapersList[1].append(self.citingPapersList[1][-1])
        except IndexError:
            self.citingPapersList[1].append(0)
        self.paperPlotInfo = {
            "id": paperID,
            "aI": self.allInfoP,
            "citList": self.citingPapersList,
        }
        if plot:
            self.paperPlotInfo["fig"] = self.plotStats(paper=True)
        if verbose > 0:
            pBLogger.info("Done!")
        return self.paperPlotInfo

    def plotStats(
        self,
        paper=False,
        author=False,
        show=False,
        save=False,
        path=".",
        markPapers=False,
        pickVal=6,
    ):
        """Plot the collected information, using matplotlib.pyplot.

        Parameters:
            paper (boolean, default False): plot statistics
                for the last analyzed paper
            author (boolean, default False): plot statistics
                for the last analyzed author
            show (boolean, default False): True to show the plots
                in a separate window (with matplotlib.pyplot.show())
            save (boolean, default False): True to save the plots into files.
            path (string): where to save the plots
            markPapers (boolean, default False): True to draw
                a vertical lines at the dates
                corresponding to a paper appearing
            pickVal (float, default 6): the picker tolerance

        Output:
            False if paper==False and author==False,
            the matplotlib.pyplot figure containing
                the citation plot if paper==True,
            a list of matplotlib.pyplot figures containing
                the various plots if author==True
        """
        if paper and self.paperPlotInfo is not None:
            if len(self.paperPlotInfo["citList"][0]) > 0:
                pBLogger.info("Plotting for paper '%s'..." % self.paperPlotInfo["id"])
                fig, ax = plt.subplots()
                plt.plot(
                    self.paperPlotInfo["citList"][0],
                    self.paperPlotInfo["citList"][1],
                    picker=pickVal,
                )
                fig.autofmt_xdate()
                if save:
                    pdf = PdfPages(osp.join(path, self.paperPlotInfo["id"] + ".pdf"))
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                return fig
        elif author and self.authorPlotInfo is not None:
            pBLogger.info("Plotting for author '%s'..." % self.authorPlotInfo["name"])
            try:
                ymin = min(
                    int(self.authorPlotInfo["allLi"][0][0].strftime("%Y")) - 2,
                    int(self.authorPlotInfo["paLi"][0][0].strftime("%Y")) - 2,
                )
                ymax = max(
                    int(self.authorPlotInfo["allLi"][0][-1].strftime("%Y")) + 2,
                    int(self.authorPlotInfo["paLi"][0][-1].strftime("%Y")) + 2,
                )
            except:
                try:
                    ymin = int(self.authorPlotInfo["paLi"][0][0].strftime("%Y")) - 2
                    ymax = int(self.authorPlotInfo["paLi"][0][-1].strftime("%Y")) + 2
                except:
                    pBLogger.warning("No publications for this author?")
                    return False
            figs = []
            if len(self.authorPlotInfo["paLi"][0]) > 0:
                fig, ax = plt.subplots()
                plt.title("Paper number")
                plt.plot(
                    self.authorPlotInfo["paLi"][0],
                    self.authorPlotInfo["paLi"][1],
                    picker=pickVal,
                )
                fig.autofmt_xdate()
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_papers.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)

            if len(self.authorPlotInfo["paLi"][0]) > 0:
                fig, ax = plt.subplots()
                plt.title("Papers per year")
                ax.hist(
                    [int(q.strftime("%Y")) for q in self.authorPlotInfo["paLi"][0]],
                    bins=range(ymin, ymax),
                    picker=True,
                )
                ax.get_xaxis().get_major_formatter().set_useOffset(False)
                plt.xlim([ymin, ymax])
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_yearPapers.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)

            if len(self.authorPlotInfo["allLi"][0]) > 0:
                fig, ax = plt.subplots()
                plt.title("Total citations")
                plt.plot(
                    self.authorPlotInfo["allLi"][0],
                    self.authorPlotInfo["allLi"][1],
                    picker=pickVal,
                )
                fig.autofmt_xdate()
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_allCit.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)

            if len(self.authorPlotInfo["allLi"][0]) > 0:
                fig, ax = plt.subplots()
                plt.title("Citations per year")
                ax.hist(
                    [int(q.strftime("%Y")) for q in self.authorPlotInfo["allLi"][0]],
                    bins=range(ymin, ymax),
                    picker=True,
                )
                ax.get_xaxis().get_major_formatter().set_useOffset(False)
                plt.xlim([ymin, ymax])
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_yearCit.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)

            if len(self.authorPlotInfo["meanLi"][0]) > 0:
                fig, ax = plt.subplots()
                plt.title("Mean citations")
                plt.plot(
                    self.authorPlotInfo["meanLi"][0],
                    self.authorPlotInfo["meanLi"][1],
                    picker=pickVal,
                )
                fig.autofmt_xdate()
                if markPapers:
                    for q in self.authorPlotInfo["paLi"][0]:
                        plt.axvline(
                            datetime.datetime(
                                int(q.strftime("%Y")),
                                int(q.strftime("%m")),
                                int(q.strftime("%d")),
                            ),
                            color="k",
                            ls="--",
                        )
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_meanCit.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)

            if len(self.authorPlotInfo["aI"].keys()) > 0:
                fig, ax = plt.subplots()
                plt.title("Citations for each paper")
                for i, p in enumerate(self.authorPlotInfo["aI"].keys()):
                    try:
                        plt.plot(
                            self.authorPlotInfo["aI"][p]["citingPapersList"][0],
                            self.authorPlotInfo["aI"][p]["citingPapersList"][1],
                        )
                    except:
                        pBLogger.exception("Something went wrong while plotting...")
                fig.autofmt_xdate()
                if save:
                    pdf = PdfPages(
                        osp.join(path, self.authorPlotInfo["name"] + "_paperCit.pdf")
                    )
                    pdf.savefig()
                    pdf.close()
                if show:
                    plt.show()
                plt.close()
                figs.append(fig)
            return figs
        else:
            pBLogger.info("Nothing to plot...")
            return False


pBStats = InspireStatsLoader()
