import plotly.figure_factory as ff
import plotly.graph_objects as go
import plotly
import json
import os
from decimal import Decimal

dir_path = os.path.dirname(os.path.realpath(__file__))

logcolorscale = [
        [
          0,
          '#482878'
        ],
        [
          0.0001,
          '#482878'
        ],
        [
          0.001,
          '#31688e'
        ],
        [
          0.01,
          '#1f9e89'
        ],
        [
          0.1,
          '#6ece58'
        ],
        [
          1,
          '#fde725'
        ]
      ]

results = []
texts = []
devicesY = []
algorithmsX = [
  "RSA1024/ECB/OAEPWithSHA-1AndMGF1Padding", 
  "RSA2048/ECB/OAEPWithSHA-1AndMGF1Padding", 
  "RSA3072/ECB/OAEPWithSHA-1AndMGF1Padding",
  "RSA4096/ECB/OAEPWithSHA-1AndMGF1Padding",

  'RSA1024/ECB/OAEPWithSHA-224AndMGF1Padding',
  'RSA2048/ECB/OAEPWithSHA-224AndMGF1Padding',
  'RSA3072/ECB/OAEPWithSHA-224AndMGF1Padding', 
  'RSA4096/ECB/OAEPWithSHA-224AndMGF1Padding', 

  'RSA1024/ECB/OAEPWithSHA-256AndMGF1Padding',
  'RSA2048/ECB/OAEPWithSHA-256AndMGF1Padding',
  'RSA3072/ECB/OAEPWithSHA-256AndMGF1Padding',
  'RSA4096/ECB/OAEPWithSHA-256AndMGF1Padding',

  'RSA2048/ECB/OAEPWithSHA-384AndMGF1Padding',
  'RSA3072/ECB/OAEPWithSHA-384AndMGF1Padding',
  'RSA4096/ECB/OAEPWithSHA-384AndMGF1Padding',

  'RSA2048/ECB/OAEPWithSHA-512AndMGF1Padding', 
  'RSA3072/ECB/OAEPWithSHA-512AndMGF1Padding', 
  'RSA4096/ECB/OAEPWithSHA-512AndMGF1Padding']

benchmarkNames = [
  'benchmarkRsa1024EcbOaepSHA1AndMgf1Padding', 
  'benchmarkRsa2048EcbOaepSHA1AndMgf1Padding', 
  'benchmarkRsa3072EcbOaepSHA1AndMgf1Padding',
  'benchmarkRsa4096EcbOaepSHA1AndMgf1Padding',
  'benchmarkRsa1024EcbOaepSHA224AndMgf1Padding',
  'benchmarkRsa2048EcbOaepSHA224AndMgf1Padding',
  'benchmarkRsa3072EcbOaepSHA224AndMgf1Padding', 
  'benchmarkRsa4096EcbOaepSHA224AndMgf1Padding', 
  'benchmarkRsa512EcbOaepSHA256AndMgf1Padding',
  'benchmarkRsa1024EcbOaepSHA256AndMgf1Padding',
  'benchmarkRsa2048EcbOaepSHA256AndMgf1Padding',
  'benchmarkRsa3072EcbOaepSHA256AndMgf1Padding',
  'benchmarkRsa4096EcbOaepSHA256AndMgf1Padding',
  'benchmarkRsa512EcbOaepSHA384AndMgf1Padding', 
  'benchmarkRsa1024EcbOaepSHA384AndMgf1Padding', 
  'benchmarkRsa2048EcbOaepSHA384AndMgf1Padding',
  'benchmarkRsa3072EcbOaepSHA384AndMgf1Padding',
  'benchmarkRsa4096EcbOaepSHA384AndMgf1Padding',
  'benchmarkRsa512EcbOaepSHA512AndMgf1Padding',
  'benchmarkRsa1024EcbOaepSHA512AndMgf1Padding',
  'benchmarkRsa2048EcbOaepSHA512AndMgf1Padding', 
  'benchmarkRsa3072EcbOaepSHA512AndMgf1Padding', 
  'benchmarkRsa4096EcbOaepSHA512AndMgf1Padding']

benchmarkClassName = 'cz.vutbr.benchmark.AsymmetricDecryptionBenchmark'

for device in os.listdir("{}/results".format(dir_path)):
    filePath = "{}/results/{}/result.json".format(dir_path, device)
    content = open(filePath, "r").read()
    devicesY.append(device)
    decoded = json.loads(content)
    benchmarkResults = decoded['benchmarks']
    result = []
    text = []
    for benchmarkName in benchmarkNames:
      for benchmark in benchmarkResults:
          if benchmark['className'] == benchmarkClassName and benchmark['name'] == benchmarkName:
              metric = benchmark['metrics']['timeNs']['median']
              result.append(metric/1000000)
              text.append('%.2f' % Decimal(metric/1000000))
              break
    results.append(result)
    texts.append(text)

fig = ff.create_annotated_heatmap(results, x=algorithmsX, y=devicesY, annotation_text=texts, colorscale=logcolorscale, showscale = True)
fig.update_layout(
    xaxis_title="Algorithm",
    yaxis_title="Device",
    autosize=False,
    width=1280,
    height=720,
)
fig['layout']['xaxis']['side'] = 'bottom'
fig['data'][0]['colorbar'].tickformat = ",.1"
fig['data'][0]['colorbar'].ticksuffix = " ms"
fig.write_image("{}/pdf/asymmetric_key_decryption_RSA_ECB_OAEPWithSHAandMGF1Padding.pdf".format(dir_path))