"""
Generates default_varmap.py.

Needs to be run only once. Needs to see a matplotlib repository.

"""

from codemend import relative_path

from element import get_pyplot_funcs
from element_extract import extract_varmap_elems

if __name__ == '__main__':
  with open(relative_path('lib/matplotlib/pyplot.py')) as reader:
    pyplot_src = reader.read()

  with open(relative_path('lib/matplotlib/pylab.py')) as reader:
    pylab_src = reader.read()

  var_map1, _ = extract_varmap_elems(pyplot_src, False, True)
  print len(var_map1)

  var_map2, _ = extract_varmap_elems(pylab_src, False, True)
  print len(var_map2)

  lines = []
  pyplot_funcs = get_pyplot_funcs()
  for f in pyplot_funcs:
    lines.append('from matplotlib.pyplot import %s'%f)

  fake_pyplot_src = '\n'.join(lines)
  var_map3, _ = extract_varmap_elems(fake_pyplot_src, False, True)
  print len(var_map3)

  var_map = dict(var_map1.items() + var_map2.items() + var_map3.items())
  print len(var_map)

  var_map['plt'] = 'matplotlib.pyplot'
  var_map['PLT'] = 'matplotlib.pyplot'

  with open(relative_path('models/default_varmap.py'), 'w') as writer:
    writer.write('""" Automatically generated by default_varmap_gen.py. """\n')
    super_long_line = 'default_var_map = %s\n'%repr(var_map)
    super_long_line = super_long_line.replace(', ', ',\n    ')
    writer.write(super_long_line)
