from __future__ import print_function,division

# Program to fit the demand curve of Newman and Ferrario (2019)
#
# Accepts input in the form of a CSV file with price in the first
# column followed by any number of columns containing numbers of
# responses for each price
#
# Written by Mark Newman, August 20, 2019
# This program may be freely distributed and/or modified

from sys import stderr,argv
from csv import reader
from numpy import array,linspace,mean,min,max,argmax,log,exp
from scipy.optimize import curve_fit
from pylab import figure,subplot,subplots_adjust,plot,savefig
from pylab import xscale,yscale,xlabel,ylabel,close

# Image format for output files: possible types include "png", "pdf",
# "eps", and "jpg"
image_format = "png"

# Fix input to work with either Python 2 or 3
try:
    input = raw_input
except NameError:
    pass

# Log of demand curve
def logQ(P,Q0,P0,a,b):
    return log(Q0) - (a/b)*log(1+(P/P0)**b)

# Get the filename
if len(argv)>1: filebase = argv[1]
else: filebase = input("Enter the name of the CSV file: ")
if ".csv" in filebase:
    endstr = filebase.find(".csv")
    filebase = filebase[:endstr]   # Remove ".csv"
if ".CSV" in filebase:
    endstr = filebase.find(".CSV")
    filebase = filebase[:endstr]   # Remove ".CSV"
filename = filebase + ".csv"
try:
    f = open(filename,"r")
except:
    print("Unable to open file",filename,file=stderr)
    exit(1)

# Read the CSV file
freader = reader(f)
rawdata = []
for row in freader: rawdata.append(row)
f.close()

# Put the data in an array and count the number of curves
data = array(rawdata,float)
nprices = data.shape[0]       # Number of price points
ncurves = data.shape[1] - 1   # Number of curves
print("Data read from file",filename)
print("  Price points:",nprices)
print("  Curves:      ",ncurves)
print()

# Open the parameter output file
outf = open(filebase+"_params.csv","w")
print('CURVE,Q0,P0,a,b,Rmax,Pmax,"Normalized Pmax"',file=outf)

# Perform the fits
for r in range(1,ncurves+1):

    print("Analyzing curve",r)

    # Pick out the non-zero data points
    Pdata = []
    Rdata = []
    for k in range(nprices):
        P = data[k,0]
        R = data[k,r]
        if R>0.0:
            Pdata.append(P)
            Rdata.append(R)
    Pdata = array(Pdata)
    Rdata = array(Rdata)
    Qdata = Rdata/Pdata

    # Compute the ranges
    Plo = min(Pdata)
    Phi = max(Pdata)
    Qlo = min(Qdata)
    Qhi = max(Qdata)

    # Initialize the fit parameters
    if len(argv)>=6:
        Q0 = float(argv[2])
        P0 = float(argv[3])
        a = float(argv[4])
        b = float(argv[5])
    else:
        R0 = argmax(Rdata)
        P0 = Pdata[R0]
        Q0 = max(Qdata)
        a = 2.0
        b = 1.0

    param_init = [Q0,P0,a,b]
    param_max = 10000.0

    # Perform the fit
    try:
        popt,pcov = curve_fit(logQ,Pdata,log(Qdata),
                              p0=param_init,bounds=(0,param_max))
    except:
        print(r,',"DID NOT CONVERGE",,,,,',sep="",file=outf)
        print("  Caution: Fit did not converge")
        continue
    Q0,P0,a,b = popt
    Pmax = P0/(a-1)**(1/b)
    Qmax = exp(logQ(Pmax,Q0,P0,a,b))
    Rmax = Pmax*Qmax
    print(r,",",Q0,",",P0,",",a,",",b,",",Rmax,",",Pmax,",",Q0*Pmax,
          sep="",file=outf)

    # Check the ranges and print warnings if parameters are out of range
    if Pmax<Plo or Pmax>Phi:
        print("  Caution: Pmax = ",float("%.2g"%Pmax)," is outside the data price range [",
              float("%.2g"%Plo)," to ",float("%.2g"%Phi),"]",sep="")
    if Q0>2*Qhi:
        print("  Caution: Q0 = ",float("%.2g"%Q0)," is substantially outside the data range [",
              float("%.2g"%Qlo)," to ",float("%.2g"%Qhi),"]",
              sep="")

    # Make figure
    figure(figsize=(10,6))
    subplots_adjust(left=0.1,right=0.95,top=0.95,bottom=0.1,wspace=0.3)

    # Plot the demand curve
    subplot(1,2,1)
    Pfit = linspace(Plo,Phi,10001)
    Qfit = exp(logQ(Pfit,Q0,P0,a,b))
    plot(Pfit,Qfit,"g-")
    plot(Pdata,Qdata,"bo")
    xscale("log")
    yscale("log")
    xlabel("Price P")
    ylabel("Consumption Q")

    # Plot the revenue curve
    subplot(1,2,2)
    Rfit = Pfit*Qfit
    plot(Pfit,Rfit,"g-")
    plot(Pdata,Rdata,"bo")
    xlabel("Price P")
    ylabel("Responses R")
    xscale("log")

    # Save and close figure
    savefig(filebase+"_"+str(r)+"."+image_format)
    close()

outf.close()