# Copyright (C) 2018-2023 Battelle Memorial Institute
# Copyright (C) 2024 Meltran, Inc
"""Load a Python dictionary of specific BES from database queries.
"""
import emthub.EMTHubConfig as EMTHubConfig
from SPARQLWrapper import SPARQLWrapper2, JSON
import time
import xml.etree.ElementTree as ET

# global constants
SPARQL = None
PREFIX = None
DELIM = ':'

def initialize_SPARQL (cfg_file=None):
  """Configure Blazegraph and initialize the SPARQL query object.

  Blazegraph must already be started. If not already configured, cfg_file must be provided.
  Writes a summary of classes and tuples to the console.

  Args:
    cfg_file (str): configuration file for Blazegraph.
  """
  global SPARQL
  if cfg_file is not None:
    EMTHubConfig.ConfigFromJsonFile (cfg_file)

  SPARQL = SPARQLWrapper2(EMTHubConfig.blazegraph_url)
  SPARQL.setReturnFormat(JSON)

def build_query (prefix, base, fid):
  if fid is not None:
    idx = base.find('WHERE {') + 8
    return prefix + '\n' + base[:idx] + """ VALUES ?fdrid {{"{:s}"}}\n""".format (fid) + base[idx:]
  return prefix + '\n' + base

def build_dict (ret):
  dict = {}
  vars = ret.variables
  keyfld = 'id'
  vars.remove(keyfld)
  for b in ret.bindings:
    key = b[keyfld].value
    row = {}
    for fld in vars:
      row[fld] = b[fld].value
    dict[key] = row
  return dict

def query_for_values (tbl, fid):
  keyflds = tbl['keyfld'].split(':')
  query = build_query (PREFIX, tbl['sparql'], fid)
  SPARQL.setQuery (query)
  ret = SPARQL.query()
  vars = ret.variables
  for akey in keyflds:
    vars.remove (akey)
  tbl['columns'] = vars
  for b in ret.bindings:
    row = {}
    key = b[keyflds[0]].value
    for i in range(1, len(keyflds)):
      key = key + DELIM + b[keyflds[i]].value
    for fld in vars:
      if fld not in b:
        row[fld] = ''
      else:
        if fld in ['name', 'phases', 'bus', 'state', 'conn', 'fdrid', 'bus1', 'bus2', 'lname', 'spacing']:
          row[fld] = b[fld].value
        else:
          try:
            row[fld] = int(b[fld].value)
          except ValueError:
            try:
              row[fld] = float(b[fld].value)
            except ValueError:
              row[fld] = b[fld].value

    tbl['vals'][key] = row

def summarize_feeder_dict (dict):
  print ('Query ID                       Key Field              Nrows Columns')
  lst = sorted(dict.keys())
  for key in lst:
    q = dict[key]
    print ('{:30s} {:21s} {:6d} {:s}'.format(key, str(q['keyfld']), len(q['vals']), str(q['columns'])))

def summarize_bes_dict(dict):
  summarize_feeder_dict(dict)

def list_feeders (dict):
  print ('Feeder Name          FID')
  SPARQL.setQuery (PREFIX + dict['DistFeeder']['sparql'])
  ret = SPARQL.query()
  for b in ret.bindings:
    print ('{:20s} {:s}'.format (b['feeder'].value, b['fid'].value))

def list_dict_table(dict, tag):
  tbl = dict[tag]
  print ('\n{:s}: key,{:s}'.format(tag, str(tbl['columns'])))
  for key, row in tbl['vals'].items():
    print ('{:s},{:s}'.format (key, ','.join(str(row[c]) for c in tbl['columns'])))

def load_feeder (dict, fid, bTime=True):
  for key in ['DistSolar', 'DistStorage', 'DistLoad', 'DistCapacitor', 'DistLinesSpacingZ', 'DistSubstation', 'DistBaseVoltage', 'DistFeeder',
              'DistBreaker', 'DistDisconnector', 'DistFuse', 'DistJumper', 'DistLoadBreakSwitch', 'DistRecloser', 'DistSectionaliser',
              'DistOverheadWire', 'DistConcentricNeutralCable', 'DistLineSpacing',
              'DistXfmrTank', 'DistXfmrBank', 'DistXfmrCodeRating', 'DistXfmrCodeNLTest', 'DistXfmrCodeSCTest',
              'DistCoordinates', 'DistRegulatorBanked', 'DistRegulatorTanked',
              'DistPowerXfmrCore', 'DistPowerXfmrMesh', 'DistSeriesCompensator', 'DistSyncMachine',
              'DistPhaseMatrix', 'DistSequenceMatrix', 'DistLinesCodeZ', 'DistLinesInstanceZ', 'DistTapeShieldCable',
              'DistBus', 'DistPowerXfmrWinding']:

    start_time = time.time()
    query_for_values (dict[key], fid)
    if bTime:
      print ('Running {:30s} took {:6.3f} s'.format (key, time.time() - start_time))
  # remove all but fid from the list of feeders
  match_fid = fid
  delete = [key for key in dict['DistFeeder']['vals'] if dict['DistFeeder']['vals'][key]['fid'] != match_fid]
  for key in delete: 
    del dict['DistFeeder']['vals'][key]

def load_bes (dict, sysid, bTime=True):
  for key in ['BESContainer', 'BESBus', 'BESBusXY', 'BESBaseVoltage', 'BESCompSeries', 'BESCompShunt', 'BESLoad', 'BESLine',
              'BESMachine', 'BESPowerXfmrWinding', 'BESPowerXfmrCore', 'BESPowerXfmrMesh',
              'BESExciter', 'BESGovernor', 'BESStabilizer', 'BESWind', 'BESSolar',
              'BESWeccREGCA', 'BESWeccREECA', 'BESWeccREPCA', 'BESWeccWTGTA', 'BESWeccWTGAA']: # 'BESCountPowerXfmrWindings'
    start_time = time.time()
    query_for_values (dict[key], sysid)
    if bTime:
      print ('Running {:30s} took {:6.3f} s'.format (key, time.time() - start_time))
  # remove all but fid from the list of feeders
#  match_fid = sysid
#  delete = [key for key in dict['DistFeeder']['vals'] if dict['DistFeeder']['vals'][key]['fid'] != match_fid]
#  for key in delete: 
#    del dict['DistFeeder']['vals'][key]

def load_feeder_dict (xml_file, fid, bTime=True, keyDelimiter=':', cfg_file=None):
  global PREFIX
  global DELIM
  DELIM = keyDelimiter
  initialize_SPARQL (cfg_file)
  # read the queries into dict
  tree = ET.parse(xml_file)
  root = tree.getroot()
  nsCIM = root.find('nsCIM').text.strip()
  nsRDF = root.find('nsRDF').text.strip()
  PREFIX = """PREFIX r: <{:s}>\nPREFIX c: <{:s}>""".format (nsRDF, nsCIM)
  dict = {}
  for query in root.findall('query'):
    qid = query.find('id').text.strip()
    dict[qid] = {}
    dict[qid]['keyfld'] = query.find('keyfld').text
    dict[qid]['sparql'] = query.find('value').text.strip()
    dict[qid]['columns'] = []
    dict[qid]['vals'] = {}

  load_feeder (dict, fid, bTime=False)
  return dict

def load_bes_dict (xml_file, sysid, bTime=True, keyDelimiter=':', cfg_file=None):
  """Query a BES network model from the Blazegraph database.

  Blazegraph must already be started. If not already configured, cfg_file must be provided.

  Args:
    xml_file (str): definition of the SPARQL queries, should be qbes.xml
    sysid (str): mRID (no underscore) for the BES ConnectivityNodeContainer
    bTime (bool): optionally time the query performance
    keyDelimiter (str): delimiter for separating composite keys
    cfg_file (str): configuration file for Blazegraph

  Returns:
    dict: a dictionary of the BES network model data
  """
  global PREFIX
  global DELIM
  DELIM = keyDelimiter
  initialize_SPARQL (cfg_file)
  # read the queries into dict
  tree = ET.parse(xml_file)
  root = tree.getroot()
  nsCIM = root.find('nsCIM').text.strip()
  nsRDF = root.find('nsRDF').text.strip()
  PREFIX = """PREFIX r: <{:s}>\nPREFIX c: <{:s}>""".format (nsRDF, nsCIM)
  dict = {}
  for query in root.findall('query'):
    qid = query.find('id').text.strip()
    dict[qid] = {}
    dict[qid]['keyfld'] = query.find('keyfld').text
    dict[qid]['sparql'] = query.find('value').text.strip()
    dict[qid]['columns'] = []
    dict[qid]['vals'] = {}

  load_bes (dict, sysid, bTime=False)
  return dict

