#!/usr/bin/env python3
import os
import psycopg2
import psycopg2.extras

import pandas as pd
import matplotlib.pyplot as plt

#    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
#        cursor.execute(sql)
#        records = cursor.fetchall()
#        pprint.pprint(records)

# INSERT ckient scaling
script='insert'
xcol='clients'
sql="""
SELECT
  testset.info,
  clients AS %s,
  round(avg(tps)) AS tps
FROM tests
  JOIN server ON (tests.server=server.server)
  JOIN testset ON (tests.set=testset.set AND tests.server=testset.server)
WHERE
  tests.script='%s' AND
  testset.category='reference'
GROUP BY server_info,testset.info,tests.server,tests.set,tests.script,clients
ORDER BY server_info,testset.info,tests.server,tests.set,tests.script,clients
;""" % (xcol,script);

# nobranch look oriented by scale_percentage
script='nobranch';
xcol='scale_percentage'
grouping='info'
ycol='tps'
sql="""
SELECT
  testset.info,
  round(100.0 * scale * 16 * 1024*1024 / (1024*1024*1024) / 
        -- TODO Remove this hack once server memory is tied to tesxcor set
        CASE WHEN tests.server='rising' and tests.set<14 THEN server_mem_gb / 2 ELSE server_mem_gb END)
        AS %s,
  round(avg(tps)) AS tps
FROM tests
  JOIN server ON (tests.server=server.server)
  JOIN testset ON (tests.set=testset.set AND tests.server=testset.server)
WHERE
  tests.script='%s' AND
  testset.category='reference'
  AND server.server in ('rising','crystal')
GROUP BY server_info,testset.info,tests.server,server_mem_gb,tests.set,tests.script,scale
ORDER BY server_info,testset.info,tests.server,server_mem_gb,tests.set,tests.script,scale
;""" % (xcol,script);

# nobranch drawdown for a fixed scale oriented by clients grouped by client limit
script='nobranch';
xcol='clients'
server='rising'
server_label = server
grouping='client_limit'
ycol='drawdown'
sql="""
SELECT
--  tests.set,
--  scale,
  client_limit,
  clients AS "%s",
  CASE WHEN avg(tps)>0 AND
      ((avg(tps)/clients - client_limit) / client_limit) < -0.04
    THEN round(100*(avg(tps)/clients - client_limit) / client_limit)
    ELSE 0 END AS drawdown
FROM tests
  JOIN server ON (tests.server=server.server)
  JOIN testset ON (tests.set=testset.set AND tests.server=testset.server)
WHERE
  client_limit > 0 AND
  tests.script='%s' AND
  tests.server='%s' AND
  testset.category='reference' AND
  scale=34250
GROUP BY tests.set,script,scale,clients,client_limit
ORDER BY tests.set,script,scale,clients,client_limit
;""" % (xcol,script,server);

server='rising'
server_label = server
xcol='clients'
grouping='info'
ycol='tps'
script='select'
# TODO add server name from lookup and use that for chart title
sql="""
SELECT
--  server_info,
  testset.info,
  clients AS %s,round(avg(tps)) AS tps
FROM tests
  JOIN server ON (tests.server=server.server)
  JOIN testset ON (tests.set=testset.set AND tests.server=testset.server)
WHERE
  tests.script='%s' AND
  tests.server='%s' AND
  testset.category='reference'
GROUP BY server_info,testset.info,tests.server,tests.set,tests.script,clients
ORDER BY server_info,testset.info,tests.server,tests.set,tests.script,clients
;""" % (xcol,script,server);

def main():
    conn_string = "host='localhost' dbname='results' user='gsmith' password='secret'"
    print("Connecting to database\n	->%s" % (conn_string))
    conn = psycopg2.connect(conn_string)

    try:
        print(sql)
        df = pd.read_sql_query(sql, conn)
        return df
    finally:
        conn.close()

if __name__ == "__main__":
    # TODO Create this directory if it doesn't exist
    base="images"
    df=main()
    df.set_index(xcol, inplace=True)
    g=df.groupby(grouping)

    #TODO this only prints an object reference, not useful
    #    print(g)

    # Disable the warning generated by the rename    
    pd.options.mode.chained_assignment = None  # default='warn'

    fig=plt.figure();
    ax=fig.add_subplot(1,1,1)
    
    for k,v in g:
        print("Processing",k)
        v.rename(columns={ycol:k},inplace=True)

        # TODO Added this for drawdown, not sure what it does to earlier graphs
        v.pop(grouping)

        print(v)
        # TODO title needs to be server_label on first SQL, then script on second/third
        ax=v.plot(grid=True,title=script,figsize=(8,6),ax=ax)
        ax.set_ylabel(ycol)
        ax.xaxis.grid(True, which="minor")

        # Overwrite this for now to get all 3 saved to one file
        k=server
        
        unslashed=k.replace("/","-")
        fn=os.path.join(base,unslashed)

        plt.savefig(fn,dpi=80)  # gives 640x480 figures
#        plt.close()
        print("saved to '%s.png'" % fn)

    # Reserve legend order for when later sets tend to be faster
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(reversed(handles), reversed(labels), loc='best')

    plt.show()
