#!/usr/bin/env python

# See LICENSE for license details.

import os
import os.path
from subprocess import Popen
import argparse
import shutil
import re
import csv 
import numpy as np
import math

parser = argparse.ArgumentParser(
  description = 'Run PrimeTime PX for each sample generated by replay-sample.py')
parser.add_argument("-s", "--sample", dest="sample", type=str,
                    help='sample file', required=True)
parser.add_argument("-d", "--design", dest="design", type=str,
                    help='design name', required=True)
parser.add_argument("-m", "--make", dest="make", type=str, nargs='*',
                    help='make command', required=True)
parser.add_argument("--output-dir", dest="output_dir", type=str,
                    help='output directory for vpd, power', required=True)
parser.add_argument("--trace-dir", dest="trace_dir", type=str,
                    help='PLSI TRACE directory', required=True)
parser.add_argument("--obj-dir", dest="obj_dir", type=str,
                    help='PLSI OBJ directory', required=True)
args = parser.parse_args()

""" Read Sample """
num = 0
with open(args.sample) as f:
  for line in f:
    tokens = line.split(" ")
    head = tokens[0]
    if head == '1':
      assert tokens[1] == 'cycle:'
      num += 1

prefix = os.path.basename(os.path.splitext(args.sample)[0])
if not os.path.exists(args.output_dir):
  os.makedirs(args.output_dir)
if not os.path.exists(args.trace_dir):
  os.makedirs(args.trace_dir)

ids = range(num)
for k in xrange(0, num, 10):
  ps = list()

  for i in ids[k:k+10]:
    """ Copy vpd """
    shutil.copy("%s/%s-replay-%d.vpd" % (args.output_dir, prefix, i),
               "%s/%s-replay-%d.vpd" % (args.trace_dir, prefix, i))

    """ Run PrimeTime PX """
    cmd = ["make"] + args.make + \
          ["SAMPLE=%s/%s-replay-%d.sample" % (args.output_dir, prefix, i)]
    ps.append(Popen(cmd, stdout=open(os.devnull, 'wb')))

  while any(p.poll() == None for p in ps):
    pass

  assert all(p.poll() == 0 for p in ps)

""" Read report file """
modules = list()
sample_pwr = dict()

for i in xrange(num):
  report_filename = "%s/pt-power/%s-replay-%d/synopsys-pt-workdir/reports/%s_report_power.report" % (
    args.obj_dir, prefix, i, args.design)
  with open(report_filename) as f:
    found = False
    for line in f:
      tokens = line.split()
      if not found:
        found = len(tokens) > 0 and tokens[0] == 'Hierarchy'
      elif found and len(tokens) >= 6:
        module = ' '.join(tokens[:2]) if len(tokens) > 6 else tokens[0]
        int_pwr = tokens[-5]
        switch_pwr = tokens[-4]
        leak_pwr = tokens[-3]
        total_pwr = tokens[-2]
        percent = tokens[-1]
        if not 'clk_gate' in module:
          if not module in sample_pwr:
            modules.append(module)
            sample_pwr[module] = list()
          sample_pwr[module].append('0.0' if total_pwr == 'N/A' else total_pwr)

""" Dump power """
csv_filename = "%s/%s-pwr.csv" % (args.output_dir, prefix)
print "[strober] Dump Power at", csv_filename
with open(csv_filename, "w") as f:
  writer = csv.writer(f)
  writer.writerow(["Modules"] + ["Sample %d (mW)" % i for i in xrange(num)] + [
    "Average (mW)", "95% error", "99% error"])
  for m in modules:
    arr = np.array([1000.0 * float(x) for x in sample_pwr[m]])
    avg = np.mean(arr)
    var = np.sum(np.power(arr - avg, 2)) / (num - 1) if num > 1 else 0 # sample variance
    _95 = 1.96 * math.sqrt(var / num)
    _99 = 2.576 * math.sqrt(var / num)
    writer.writerow([m] + arr.tolist() + [avg, _95, _99])
