#!/usr/bin/env python3
# Copyright (C) 2021 Jonathan A. Poritz
# 
# This program 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.
# 
# This program 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.
#
#
#

# handle command-line arguments
import argparse
parser = argparse.ArgumentParser(description="Make a graph of books in an OER repo")
parser.add_argument('--startyear', '-s', action='store',
                    dest='start_year',
                    type=int,
                    default=2000,
                    help='Year to start graph; default: 2000')
parser.add_argument('--datafile', '-d', action='store',
                    dest='data_file',
                    default='otlsize',
                    help="file to get data/size info; default: 'otlsize'; should be in format: YYYY MM DD num")
parser.add_argument('--reponame', '-r', action='store',
                    dest='repo_name',
                    default='the Open Textbook Library',
                    help="name of the repository; default: 'Open Textbok Library'")
parser.add_argument('--topy', '-t', action='store',
                    dest='y_axis_top',
                    type=int,
                    default=0,
                    help="Value for top of graph's y-axis; default: automatic")
parser.add_argument('--lsrlr','-l', action='append',
                    dest='lsrl_range',
                    default=[],
                    help="range for plotting LSRL for some of the data; may be repeated; default: no LSRLs plotted; should be in format: '(lower,upper)' or '(lower,upper):(a,b)' in case information about the LSRL should be shown at location (a,b) in the graph's coordinate system")
parser.add_argument('--expreg','-e', action='append',
                    dest='er_range',
                    default=[],
                    help="range for plotting exponential regression for some of the data; may be repeated; default: no exponential regression plotted; should be in format: '(lower,upper)' or '(lower,upper):(a,b)' in case information about the exponential regression should be shown at location (a,b) in the graph's coordinate system")
parser.add_argument('-x', action="store_true", default=False,
                    dest='show_xs', help="show list of date values found in data file")
parser.add_argument('-y', action="store_true", default=False,
                    dest='show_ys', help="show list of repo size values found in data file")
results = parser.parse_args()

# parse and check sanity of LSRL command-line arguments
lsrlranges=[]
lsrlinfolocs=[]
for r in results.lsrl_range:
  cloc=r.find(':')
  if cloc>=0:
    a=r[:cloc]
    b=r[cloc+1:]
    assert b[0]=='(' and b[-1]==')', "malformed LSRL info location"
    lower_end=b.find(',')
    assert lower_end>0, "malformed LSRL info location"
    xcoord=int(b[1:lower_end])
    ycoord=int(b[lower_end+1:-1])
    lsrlinfolocs.append([xcoord,ycoord])
  else:
    a=r
    lsrlinfolocs.append([])
  assert a[0]=='(' and a[-1]==')', "malformed LSRL range"
  lower_end=a.find(',')
  assert lower_end>0, "malformed LSRL range"
  lower=int(a[1:lower_end])
  upper=int(a[lower_end+1:-1])
  assert lower<upper, "malformed LSRL range"
  lsrlranges.append([lower,upper])

# parse and check sanity of expreg command-line arguments
erranges=[]
erinfolocs=[]
for r in results.er_range:
  cloc=r.find(':')
  if cloc>=0:
    a=r[:cloc]
    b=r[cloc+1:]
    assert b[0]=='(' and b[-1]==')', "malformed exponential regression info location"
    lower_end=b.find(',')
    assert lower_end>0, "malformed exponential regression info location"
    xcoord=int(b[1:lower_end])
    ycoord=int(b[lower_end+1:-1])
    erinfolocs.append([xcoord,ycoord])
  else:
    a=r
    erinfolocs.append([])
  assert a[0]=='(' and a[-1]==')', "malformed exponential regression range"
  lower_end=a.find(',')
  assert lower_end>0, "malformed exponential regression range"
  lower=int(a[1:lower_end])
  upper=int(a[lower_end+1:-1])
  assert lower<upper, "malformed exponential regression range"
  erranges.append([lower,upper])

# read data file and turn dates into days offset from 1 Jan of whatever is the
# starting year.  Note the data file should consist only of lines of the form
# YYYY MM DD <number>
# where date specs are numbers (use leading 0s to get to two digits) and
# <number> can be any number of digits
monthsizes=[31,28,31,30,31,30,31,31,30,31,30,31]
lymonthsizes=[31,29,31,30,31,30,31,31,30,31,30,31]
x=[]
y=[]
fh = open(results.data_file, "r")
s=fh.readline()
while(len(s)):
  year=int(s[:4])
  month=int(s[5:7])-1
  day=int(s[8:10])-1
  if results.start_year<=year:
    offset=sum([365-(yr%4==0) for yr in range(results.start_year,year)])
  else:
    offset=-sum([365-(yr%4==0) for yr in range(year,results.start_year)])
  if year%4==0:
    offset += sum(lymonthsizes[0:month])
  else:
    offset += sum(monthsizes[0:month])
  offset += day
  x.append(offset)
  y.append(int(s[11:]))
  s=fh.readline()
if results.show_xs:
  print("x:",x)
if results.show_ys:
  print("y:",y)

# do the basic scatterplot
import matplotlib.pyplot as plt
plt.plot(x,y,'r.')
plt.rc_context({'axes.autolimit_mode': 'round_numbers'})
if results.y_axis_top:
  plt.ylim(bottom=0,top=results.y_axis_top)
else:
  plt.ylim(bottom=0)
plt.xlabel("Days since 1 January "+str(results.start_year))
plt.ylabel("Books")
plt.title("Books in "+results.repo_name+" over time")

# do any requested linear regressions
import numpy as np
for (r,s) in zip(lsrlranges,lsrlinfolocs):
  ax=[]
  ay=[]
  for (i,j) in zip(x,y):
    if i>=r[0] and i<=r[1]:
      ax.append(i)
      ay.append(j)
  m,b=np.polyfit(ax,ay,1)
  ry=[r[0]*m+b,r[1]*m+b]
  plt.plot(r,ry,'g-')
  print("For the regression line in x-range ("+str(r[0])+","+str(r[1])+"):")
  print("  m="+str(m)[:5]+"; meaning: a new book every "+str(int(round((1/m)*10))/10)+" days")
  print("  b="+str(b)[:5])
  print("  r="+str(np.corrcoef(ax,ay)[0][1])[1:5])
  if len(s):
    plt.text(s[0],s[1],"m="+str(m)[:5]+"; b="+str(b)[:5]+"; r="+str(np.corrcoef(ax,ay)[0][1])[1:5]+"\na new book every "+str(int(round((1/m)*10))/10)+" days",color='g')

# do any requested exponential regressions
for (r,s) in zip(erranges,erinfolocs):
  bx=[]
  by=[]
  for (i,j) in zip(x,y):
    if i>=r[0] and i<=r[1]:
      bx.append(i)
      by.append(j)
  pf=np.polynomial.Polynomial.fit(x, np.log(y), deg=1)
  lx=np.linspace(r[0],r[1])
  ly=np.exp(pf(lx))
  plt.plot(lx,ly,'g-')
  C=np.exp(pf(0))
  a=np.exp(pf(1))/C
  print("For the exponential regression y = Ca^x in x-range ("+str(r[0])+","+str(r[1])+"):")
  print("  C="+str(C)[:5]+" and a="+str(a)[:9]+", meaning: # of books doubles in "+str(int(round(100*np.log(2)/np.log(a)))/100)+" days/"+str(int(round(100*np.log(2)/(365.25*np.log(a))))/100)+" years")
  if len(s):
    plt.text(s[0],s[1],"C="+str(C)[:5]+"; a="+str(a)[:9]+"; # books doubles\nevery "+str(int(round(100*np.log(2)/np.log(a)))/100)+" days/"+str(int(round(100*np.log(2)/(365.25*np.log(a))))/100)+" years",color='g')

# finally, show the plot
plt.show()