#!/usr/bin/env python3

# Compute analytical formulas for the bytes used per epoch by the
# various modes

import os
import sympy

from bwparser import BandwidthParser


# sensible defaults
bandwidth_file = os.getenv('BW_FILE')
network_size = 0
if os.getenv('BW_ALGO') != "komlo":
    bw_parser = BandwidthParser(bw_file=bandwidth_file)
    network_size = bw_parser.get_relay_num()
else:
    # keep the original assumption
    network_size = 6500

A, R_B, R_N, R, logR, C_B, C_N, C, gamma, circ, P_Delta, \
DirAuthConsensusMsg, DirAuthENDIVEDiffMsg, DirAuthGetConsensusMsg, \
DirAuthGetENDIVEDiffMsg, DirAuthUploadDescMsg, DirAuthENDIVEMsg, \
DirAuthGetENDIVEMsg, DirAuthUploadDescMsg, RelayConsensusMsg, \
RelayDescMsg, RelayGetConsensusMsg, RelayGetDescMsg, \
SinglePassCreateCircuitMsgNotLast, SinglePassCreateCircuitMsgLast, \
SinglePassCreatedCircuitCellLast, SinglePassCreatedCircuitCellMiddle, \
SinglePassCreatedCircuitCellFirst, TelescopingCreateCircuitMsg, \
TelescopingCreatedCircuitCell, TelescopingExtendCircuitCell, \
TelescopingExtendedCircuitCell, VanillaCreateCircuitMsg, \
VanillaCreatedCircuitCell, VanillaExtendCircuitCell, \
VanillaExtendedCircuitCell, DirAuthGetConsensusDiffMsg, \
DirAuthConsensusDiffMsg, RelayGetConsensusDiffMsg, \
RelayConsensusDiffMsg \
= sympy.symbols("""
A, R_B, R_N, R, logR, C_B, C_N, C, gamma, circ, P_Delta,
DirAuthConsensusMsg, DirAuthENDIVEDiffMsg, DirAuthGetConsensusMsg,
DirAuthGetENDIVEDiffMsg, DirAuthUploadDescMsg, DirAuthENDIVEMsg,
DirAuthGetENDIVEMsg, DirAuthUploadDescMsg, RelayConsensusMsg,
RelayDescMsg, RelayGetConsensusMsg, RelayGetDescMsg,
SinglePassCreateCircuitMsgNotLast, SinglePassCreateCircuitMsgLast,
SinglePassCreatedCircuitCellLast, SinglePassCreatedCircuitCellMiddle,
SinglePassCreatedCircuitCellFirst, TelescopingCreateCircuitMsg,
TelescopingCreatedCircuitCell, TelescopingExtendCircuitCell,
TelescopingExtendedCircuitCell, VanillaCreateCircuitMsg,
VanillaCreatedCircuitCell, VanillaExtendCircuitCell,
VanillaExtendedCircuitCell, DirAuthGetConsensusDiffMsg,
DirAuthConsensusDiffMsg, RelayGetConsensusDiffMsg,
RelayConsensusDiffMsg
""")

globalsubs = [
    (A , 9),
    (R_N , R - R_B),
    (R_B , 0.010 * R),
    (C_N , C - C_B),
    (C_B , 0.16 * C),
    (circ , gamma * C),
    (gamma , 8.9),
    (C , 2500000*R/network_size),
    (P_Delta, 0.019),
]

# The actual sizes in bytes of each message type were logged by
# uncommenting this line in network.py:
    # logging.info("%s size %d", type(self).__name__, sz)

singlepass_merkle_subs = [
    (DirAuthConsensusMsg, 877),
    (DirAuthGetConsensusMsg, 41),
    (DirAuthGetENDIVEMsg, 38),
    (DirAuthGetENDIVEDiffMsg, 42),
    (DirAuthENDIVEDiffMsg, (P_Delta * DirAuthENDIVEMsg).subs(globalsubs)),
    (DirAuthENDIVEMsg, 274 * R),
    (DirAuthUploadDescMsg, 425),
    (RelayConsensusMsg, 873),
    (RelayDescMsg, 415),
    (RelayGetConsensusMsg, 37),
    (RelayGetDescMsg, 32),
    (SinglePassCreateCircuitMsgLast, 187),
    (SinglePassCreateCircuitMsgNotLast, 239),
    (SinglePassCreatedCircuitCellFirst, 1426+82*logR),
    (SinglePassCreatedCircuitCellMiddle, 903+41*logR),
    (SinglePassCreatedCircuitCellLast, 190),
]

singlepass_threshsig_subs = [
    (DirAuthConsensusMsg, 789),
    (DirAuthGetConsensusMsg, 41),
    (DirAuthGetENDIVEMsg, 38),
    (DirAuthGetENDIVEDiffMsg, 42),
    (DirAuthENDIVEDiffMsg, DirAuthENDIVEMsg),
    (DirAuthENDIVEMsg, 348*R),
    (DirAuthUploadDescMsg, 425),
    (RelayConsensusMsg, 784),
    (RelayDescMsg, 415),
    (RelayGetConsensusMsg, 37),
    (RelayGetDescMsg, 32),
    (SinglePassCreateCircuitMsgLast, 187),
    (SinglePassCreateCircuitMsgNotLast, 239),
    (SinglePassCreatedCircuitCellFirst, 1554),
    (SinglePassCreatedCircuitCellMiddle, 969),
    (SinglePassCreatedCircuitCellLast, 190),
]

telescoping_merkle_subs = [
    (DirAuthConsensusMsg, 877),
    (DirAuthGetConsensusMsg, 41),
    (DirAuthGetENDIVEMsg, 38),
    (DirAuthGetENDIVEDiffMsg, 42),
    (DirAuthENDIVEDiffMsg, (P_Delta * DirAuthENDIVEMsg).subs(globalsubs)),
    (DirAuthENDIVEMsg, 234 * R),
    (DirAuthUploadDescMsg, 372),
    (RelayConsensusMsg, 873),
    (RelayGetConsensusMsg, 37),
    (RelayGetDescMsg, 32),
    (RelayDescMsg, 362),
    (TelescopingCreateCircuitMsg, 120),
    (TelescopingCreatedCircuitCell, 179),
    (TelescopingExtendCircuitCell, 122),
    (TelescopingExtendedCircuitCell, 493+41*logR),
]

telescoping_threshsig_subs = [
    (DirAuthConsensusMsg, 789),
    (DirAuthGetConsensusMsg, 41),
    (DirAuthGetENDIVEMsg, 38),
    (DirAuthGetENDIVEDiffMsg, 42),
    (DirAuthENDIVEDiffMsg, DirAuthENDIVEMsg),
    (DirAuthENDIVEMsg, 307*R),
    (DirAuthUploadDescMsg, 372),
    (RelayConsensusMsg, 788),
    (RelayGetConsensusMsg, 37),
    (RelayGetDescMsg, 32),
    (RelayDescMsg, 362),
    (TelescopingCreateCircuitMsg, 120),
    (TelescopingCreatedCircuitCell, 179),
    (TelescopingExtendCircuitCell, 122),
    (TelescopingExtendedCircuitCell, 556),
]

vanilla_subs = [
    (DirAuthConsensusDiffMsg, (P_Delta * DirAuthConsensusMsg).subs(globalsubs)),
    (DirAuthConsensusMsg, RelayConsensusMsg),
    (DirAuthGetConsensusDiffMsg, 45),
    (DirAuthGetConsensusMsg, 41),
    (DirAuthUploadDescMsg, 372),
    (RelayConsensusDiffMsg, (P_Delta * RelayConsensusMsg).subs(globalsubs)),
    (RelayConsensusMsg, 219*R),
    (RelayGetConsensusDiffMsg, 41),
    (RelayGetConsensusMsg, 37),
    (VanillaCreateCircuitMsg, 116),
    (VanillaCreatedCircuitCell, 175),
    (VanillaExtendCircuitCell, 157),
    (VanillaExtendedCircuitCell, 176),
]

# The formulas were output by bytecounts.py

singlepass_totrelay = \
  R_N * ( DirAuthConsensusMsg + DirAuthENDIVEDiffMsg + DirAuthGetConsensusMsg + DirAuthGetENDIVEDiffMsg + A*DirAuthUploadDescMsg ) \
+ R_B * ( DirAuthConsensusMsg + DirAuthENDIVEMsg + DirAuthGetConsensusMsg + DirAuthGetENDIVEMsg + A*DirAuthUploadDescMsg ) \
+ C   * ( RelayConsensusMsg + RelayDescMsg + RelayGetConsensusMsg + RelayGetDescMsg ) \
+ circ * ( 3*SinglePassCreateCircuitMsgNotLast + 2*SinglePassCreateCircuitMsgLast + 2*SinglePassCreatedCircuitCellLast + 2*SinglePassCreatedCircuitCellMiddle + SinglePassCreatedCircuitCellFirst + 20 )

singlepass_totclient = \
  C * ( RelayConsensusMsg + RelayDescMsg + RelayGetConsensusMsg + RelayGetDescMsg ) \
+ circ * ( SinglePassCreateCircuitMsgNotLast + SinglePassCreatedCircuitCellFirst + 4 )

telescoping_totrelay = \
  R_N * ( DirAuthConsensusMsg + DirAuthENDIVEDiffMsg + DirAuthGetConsensusMsg + DirAuthGetENDIVEDiffMsg + A*DirAuthUploadDescMsg ) \
+ R_B * ( DirAuthConsensusMsg + DirAuthENDIVEMsg + DirAuthGetConsensusMsg + DirAuthGetENDIVEMsg + A*DirAuthUploadDescMsg ) \
+ C * ( RelayConsensusMsg + RelayDescMsg + RelayGetConsensusMsg + RelayGetDescMsg ) \
+ circ * ( 5*TelescopingCreateCircuitMsg + 5*TelescopingCreatedCircuitCell + 4*TelescopingExtendCircuitCell + 4*TelescopingExtendedCircuitCell + 52 )

telescoping_totclient = \
  C * ( RelayConsensusMsg + RelayDescMsg + RelayGetConsensusMsg + RelayGetDescMsg ) \
+ circ * ( TelescopingCreateCircuitMsg + TelescopingCreatedCircuitCell + 2*TelescopingExtendCircuitCell + 2*TelescopingExtendedCircuitCell + 20 )

vanilla_totrelay = \
  R_N * ( DirAuthConsensusDiffMsg + DirAuthGetConsensusDiffMsg + A*DirAuthUploadDescMsg ) \
+ R_B * ( DirAuthConsensusMsg + DirAuthGetConsensusMsg + A*DirAuthUploadDescMsg ) \
+ C_N * ( RelayConsensusDiffMsg + RelayGetConsensusDiffMsg ) \
+ C_B * ( RelayConsensusMsg + RelayGetConsensusMsg ) \
+ circ * ( 5*VanillaCreateCircuitMsg + 5*VanillaCreatedCircuitCell + 4*VanillaExtendCircuitCell + 4*VanillaExtendedCircuitCell + 52 )

vanilla_totclient = \
  C_N * ( RelayConsensusDiffMsg + RelayGetConsensusDiffMsg ) \
+ C_B * ( RelayConsensusMsg + RelayGetConsensusMsg ) \
+ circ * ( VanillaCreateCircuitMsg + VanillaCreatedCircuitCell + 2*VanillaExtendCircuitCell + 2*VanillaExtendedCircuitCell + 20 )

# Copy the output into plotdats.py, replacing 'R' by 'x' and 'logR' by
# 'cail(log(x)/log(2))'

def calculate_relay_analyticals():
    relay_perclient_analyticals = {
        'singlepass_merkle': (singlepass_totrelay/C).subs(globalsubs).subs(singlepass_merkle_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'singlepass_threshsig': (singlepass_totrelay/C).subs(globalsubs).subs(singlepass_threshsig_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'telescoping_merkle':  (telescoping_totrelay/C).subs(globalsubs).subs(telescoping_merkle_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'telescoping_threshsig': (telescoping_totrelay/C).subs(globalsubs).subs(telescoping_threshsig_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'vanilla_none': (vanilla_totrelay/C).subs(globalsubs).subs(vanilla_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
    }
    return relay_perclient_analyticals

def calculate_client_analyticals():
    client_perclient_analyticals = {
        'singlepass_merkle': (singlepass_totclient/C).subs(globalsubs).subs(singlepass_merkle_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'singlepass_threshsig': (singlepass_totclient/C).subs(globalsubs).subs(singlepass_threshsig_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'telescoping_merkle':  (telescoping_totclient/C).subs(globalsubs).subs(telescoping_merkle_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'telescoping_threshsig': (telescoping_totclient/C).subs(globalsubs).subs(telescoping_threshsig_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
        'vanilla_none': (vanilla_totclient/C).subs(globalsubs).subs(vanilla_subs).simplify().replace("logR", "ceil(log(x)/log(2))").replace("R", "x"),
    }
    return client_perclient_analyticals

if __name__ == '__main__':
    print("Client analyticals: ", calculate_client_analyticals())
    print("Relay analyticals: ", calculate_relay_analyticals())
