from __future__ import absolute_import, division, print_function

# Assorted info boxes for displaying essential file information

import wxtbx.plots
import wxtbx.icons
import wx
import wx.lib
import re
import cStringIO
import os
import sys

from libtbx.utils import to_unicode, to_str

ALN_FLAGS = wx.ALL|wx.ALIGN_CENTER_VERTICAL

class InfoPanelBase(wx.MiniFrame):
  def __init__(self, *args, **kwds):
    kwds = dict(kwds)
    kwds['style'] = wx.CAPTION|wx.CLOSE_BOX|wx.RESIZE_BORDER
    if (wx.Platform == '__WXMSW__'):
      kwds['style'] |= wx.SYSTEM_MENU
    wx.MiniFrame.__init__(self, *args, **kwds)
    self.panel = wx.lib.scrolledpanel.ScrolledPanel(
      self, -1, style=wx.RAISED_BORDER)
    self.panel.SetupScrolling()
    self.panel_sizer = wx.BoxSizer(wx.VERTICAL)
    self.panel.SetSizer(self.panel_sizer)
    self.info_panel = None
    self.file_name = None
    self.Bind(wx.EVT_CLOSE, self.OnClose)
    self.Bind(wx.EVT_WINDOW_DESTROY, self.OnDestroy)

  def set_text(self, filename, text):
    '''
    Basic function for setting the title to a filename and displaying text
    '''
    self.SetTitle('Information for %s' %
                  os.path.basename(to_unicode(filename)))
    panel_text = wx.StaticText(self.panel, -1, text)
    self.panel_sizer.Add(panel_text, 0)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

  def OnClose(self, event):
    self.Destroy()

  def OnDestroy(self, event):
    pass

class ReflectionFileInfo(InfoPanelBase):
  def __init__(self, *args, **kwds):
    super(ReflectionFileInfo, self).__init__(*args, **kwds)
    self.miller_arrays = None
    self._hkl_in = None
    self._current_array = None
    box = wx.BoxSizer(wx.HORIZONTAL)
    self.panel_sizer.Add(box, 0)
    bmp = wx.StaticBitmap(self.panel, -1, wxtbx.icons.hkl_file.GetBitmap())
    box.Add(bmp, 0, ALN_FLAGS, 5)
    grid = wx.FlexGridSizer(cols=2)
    box.Add(grid, 0, ALN_FLAGS, 5)
    txt1 = wx.StaticText(self.panel, -1, "File name:")
    font = txt1.GetFont()
    font.SetWeight(wx.FONTWEIGHT_BOLD)
    txt1.SetFont(font)
    grid.Add(txt1, 0, ALN_FLAGS, 5)
    self.file_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.file_txt, 0, ALN_FLAGS, 5)
    txt3 = wx.StaticText(self.panel, -1, "Data array:")
    txt3.SetFont(font)
    grid.Add(txt3, 0, ALN_FLAGS, 5)
    self.array_choice = wx.Choice(self.panel, -1, size=(200,-1))
    grid.Add(self.array_choice, 0, ALN_FLAGS, 5)
    self.Bind(wx.EVT_CHOICE, self.OnChangeArray, self.array_choice)
    self.Centre(wx.BOTH)

  def OnChangeArray(self, event):
    array_label = self.array_choice.GetStringSelection()
    self.set_miller_array(array_label)

  def set_file(self, file_name):
    self.file_name = os.path.abspath(file_name)
    from iotbx import file_reader
    self._hkl_in = file_reader.any_file(file_name, force_type="hkl",
      raise_sorry_if_errors=True,
      raise_sorry_if_not_expected_format=True)
    self._hkl_in.check_file_type("hkl")
    self.SetTitle("Info for %s" % to_unicode(self.file_name))
    self.file_txt.SetLabel(to_unicode(self.file_name))
    self.miller_arrays = self._hkl_in.file_server.miller_arrays
    labels = [ array.info().label_string() for array in self.miller_arrays ]
    self.array_choice.SetItems(labels)
    self.array_choice.SetSelection(0)
    self.set_miller_array(labels[0])

  def set_miller_array(self, array_label):
    assert (self.file_name is not None) and (self.miller_arrays is not None)
    array = None
    for array_ in self.miller_arrays :
      if (array_.info().label_string() == array_label):
        array = array_
    assert (array is not None)
    self._current_array = array
    info_out = cStringIO.StringIO()
    array.show_comprehensive_summary(f=info_out)
    array.show_comprehensive_summary()
    info_list = info_out.getvalue().splitlines()
    if (self.info_panel is not None):
      self.panel_sizer.Detach(self.info_panel)
      self.info_panel.Destroy()
    self.info_panel = wx.Panel(self.panel, -1)
    self.panel_sizer.Add(self.info_panel, 1, wx.ALL|wx.EXPAND)
    szr = wx.BoxSizer(wx.VERTICAL)
    self.info_panel.SetSizer(szr)
    box = wx.StaticBox(self.info_panel, -1, "Array info:")
    box_szr = wx.StaticBoxSizer(box, wx.VERTICAL)
    szr.Add(box_szr, 1, wx.EXPAND|wx.ALL, 5)
    grid = wx.FlexGridSizer(cols=2)
    box_szr.Add(grid, 1, wx.EXPAND)
    for line in info_list[1:] :
      fields = line.split(":")
      label = fields[0]
      value = ":".join(fields[1:])
      txt1 = wx.StaticText(self.info_panel, -1, label + ":")
      font = txt1.GetFont()
      font.SetWeight(wx.FONTWEIGHT_BOLD)
      txt1.SetFont(font)
      grid.Add(txt1, 0, ALN_FLAGS, 5)
      txt2 = wx.StaticText(self.info_panel, -1, value)
      font2 = txt2.GetFont()
      font2.SetFamily(wx.FONTFAMILY_MODERN)
      txt2.SetFont(font2)
      grid.Add(txt2, 0, ALN_FLAGS, 5)
    if (array.is_complex_array()):
      btn = wx.Button(self.info_panel, -1, "Show map statistics")
      szr.Add(btn, 0, wx.ALL, 5)
      self.Bind(wx.EVT_BUTTON, self.OnShowMapStats, btn)
    szr.Fit(self.info_panel)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

  def OnShowMapStats(self, event):
    map_info = MapCoeffsInfo(self.GetTopLevelParent(), -1,
      "Map statistics for %s" % self._current_array.info().label_string())
    map_info.set_file(self.file_name)
    map_info.set_map_coeffs(self._current_array)
    map_info.Show()

class PDBFileInfo(InfoPanelBase):
  def __init__(self, *args, **kwds):
    super(PDBFileInfo, self).__init__(*args, **kwds)
    self._pdb_in = None
    self._hierarchy = None
    box = wx.BoxSizer(wx.HORIZONTAL)
    self.panel_sizer.Add(box, 0)
    bmp = wx.StaticBitmap(self.panel, -1, wxtbx.icons.pdb_file.GetBitmap())
    box.Add(bmp, 0, ALN_FLAGS, 5)
    grid = wx.FlexGridSizer(cols=2)
    box.Add(grid, 0, ALN_FLAGS, 5)
    txt1 = wx.StaticText(self.panel, -1, "File name:")
    font = txt1.GetFont()
    font.SetWeight(wx.FONTWEIGHT_BOLD)
    txt1.SetFont(font)
    grid.Add(txt1, 0, ALN_FLAGS, 5)
    self.file_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.file_txt, 0, ALN_FLAGS, 5)
    self.Centre(wx.BOTH)

  def set_file(self, file_name):
    self.file_name = os.path.abspath(file_name)
    from iotbx import file_reader
    import iotbx.pdb
    self._pdb_in = file_reader.any_file(file_name, force_type="pdb",
      raise_sorry_if_errors=True,
      raise_sorry_if_not_expected_format=True)
    self._pdb_in.check_file_type("pdb")
    self._hierarchy = self._pdb_in.file_object.hierarchy
    info_list = iotbx.pdb.show_file_summary(
      pdb_in=self._pdb_in.file_object,
      hierarchy=self._hierarchy)
    self.SetTitle("Info for %s" % to_unicode(self.file_name))
    self.file_txt.SetLabel(to_unicode(self.file_name))
    if (self.info_panel is not None):
      self.panel_sizer.Detach(self.info_panel)
      self.info_panel.Destroy()
    self.info_panel = wx.Panel(self.panel, -1)
    self.panel_sizer.Add(self.info_panel, 1, wx.ALL|wx.EXPAND)
    szr = wx.BoxSizer(wx.VERTICAL)
    self.info_panel.SetSizer(szr)
    box = wx.StaticBox(self.info_panel, -1, "File contents:")
    box_szr = wx.StaticBoxSizer(box, wx.VERTICAL)
    szr.Add(box_szr, 1, wx.EXPAND|wx.ALL, 5)
    grid = wx.FlexGridSizer(cols=2)
    box_szr.Add(grid, 1, wx.EXPAND)
    for label, value in info_list :
      txt1 = wx.StaticText(self.info_panel, -1, label + ":")
      font = txt1.GetFont()
      font.SetWeight(wx.FONTWEIGHT_BOLD)
      txt1.SetFont(font)
      grid.Add(txt1, 0, ALN_FLAGS, 5)
      str_value = to_str(value)
      alert = False
      if (str_value.endswith("***")):
        str_value = re.sub("\s*\*\*\*", "", str_value)
        alert = True
      txt2 = wx.StaticText(self.info_panel, -1, str_value)
      font2 = txt2.GetFont()
      font2.SetFamily(wx.FONTFAMILY_MODERN)
      if (alert):
        font2.SetWeight(wx.FONTWEIGHT_BOLD)
        txt2.SetForegroundColour((200,0,0))
        txt1.SetForegroundColour((200,0,0))
      txt2.SetFont(font2)
      grid.Add(txt2, 0, ALN_FLAGS, 5)
    if (len(self._hierarchy.models()) == 1):
      if (len(self._hierarchy.models()[0].chains()) > 1):
        btn = wx.Button(self.info_panel, -1, "B-factors by chain...")
        szr.Add(btn, 0, wx.ALL, 5)
        self.Bind(wx.EVT_BUTTON, self.OnShowChainBstats, btn)
    szr.Fit(self.info_panel)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

  def OnShowChainBstats(self, event):
    assert (self._hierarchy is not None)
    b_panel = PDBChainBisoPanel(self)
    b_panel.set_file(
      file_name=self.file_name,
      hierarchy=self._hierarchy)
    b_panel.Show()

class PDBChainBisoPanel(InfoPanelBase):
  def __init__(self, *args, **kwds):
    super(PDBChainBisoPanel, self).__init__(*args, **kwds)
    self._hierarchy = None
    box = wx.BoxSizer(wx.HORIZONTAL)
    self.panel_sizer.Add(box, 0)
    bmp = wx.StaticBitmap(self.panel, -1, wxtbx.icons.pdb_file.GetBitmap())
    box.Add(bmp, 0, ALN_FLAGS, 5)
    grid = wx.FlexGridSizer(cols=2)
    box.Add(grid, 0, ALN_FLAGS, 5)
    txt1 = wx.StaticText(self.panel, -1, "File name:")
    font = txt1.GetFont()
    font.SetWeight(wx.FONTWEIGHT_BOLD)
    txt1.SetFont(font)
    grid.Add(txt1, 0, ALN_FLAGS, 5)
    self.file_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.file_txt, 0, ALN_FLAGS, 5)
    self.Centre(wx.BOTH)

  def set_file(self, file_name, hierarchy=None):
    self.file_name = os.path.abspath(file_name)
    from scitbx.array_family import flex
    if (hierarchy is None):
      from iotbx import file_reader
      import iotbx.pdb
      pdb_in = file_reader.any_file(file_name, force_type="pdb",
        raise_sorry_if_errors=True,
        raise_sorry_if_not_expected_format=True)
      pdb_in.check_file_type("pdb")
      hierarchy = pdb_in.file_object.hierarchy
    if (len(hierarchy.models()) > 1):
      raise Sorry("Multi-MODEL PDB files not supported.")
    self._hierarchy = hierarchy
    self.SetTitle("B-factors by chain for %s" % to_unicode(self.file_name))
    self.file_txt.SetLabel(to_unicode(self.file_name))
    chain_list = wx.ListCtrl(self.panel, -1, style=wx.LC_REPORT, size=(480,160))
    chain_list.InsertColumn(0, "Chain info")
    chain_list.InsertColumn(1, "Mean B-iso (range)")
    chain_list.SetColumnWidth(0, 260)
    chain_list.SetColumnWidth(1, 200)
    for chain in hierarchy.models()[0].chains():
      n_res = len(chain.residue_groups())
      chain_atoms = chain.atoms()
      n_atoms = len(chain_atoms)
      main_conf = chain.conformers()[0]
      chain_type = "other"
      if (main_conf.is_protein()):
        chain_type = "protein"
      elif (main_conf.is_na()):
        chain_type = "nucleic acid"
      chain_info = "'%s' (%s, %d res., %d atoms)" % (chain.id, chain_type,
        n_res, n_atoms)
      b_iso = chain_atoms.extract_b()
      b_max = flex.max(b_iso)
      b_min = flex.min(b_iso)
      b_mean = flex.mean(b_iso)
      b_info = "%.2f (%.2f - %.2f)" % (b_mean, b_min, b_max)
      item = chain_list.InsertStringItem(sys.maxunicode, chain_info)
      chain_list.SetStringItem(item, 1, b_info)
    self.panel_sizer.Add(chain_list, 1, wx.EXPAND|wx.ALL, 5)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

class ImageFileInfo(InfoPanelBase):
  def __init__(self, *args, **kwds):
    super(ImageFileInfo, self).__init__(*args, **kwds)
    self._img_in = None
    box = wx.BoxSizer(wx.HORIZONTAL)
    self.panel_sizer.Add(box, 0)
    bmp = wx.StaticBitmap(self.panel, -1, wxtbx.icons.img_file.GetBitmap())
    box.Add(bmp, 0, ALN_FLAGS, 5)
    grid = wx.FlexGridSizer(cols=2)
    box.Add(grid, 0, ALN_FLAGS, 5)
    txt1 = wx.StaticText(self.panel, -1, "File name:")
    font = txt1.GetFont()
    font.SetWeight(wx.FONTWEIGHT_BOLD)
    txt1.SetFont(font)
    grid.Add(txt1, 0, ALN_FLAGS, 5)
    self.file_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.file_txt, 0, ALN_FLAGS, 5)
    self.Centre(wx.BOTH)

  def set_file(self, file_name):
    self.file_name = os.path.abspath(file_name)
    from iotbx import file_reader
    img_in = file_reader.any_file(
      file_name=file_name,
      valid_types=["img"],
      force_type="img",
      raise_sorry_if_errors=True,
      raise_sorry_if_not_expected_format=True)
    img_in.assert_file_type("img")
    self._img_in = img_in
    out = cStringIO.StringIO()
    img_in.file_object.show_header()
    img_in.file_object.show_header(out=out)
    self.SetTitle("Info for %s" % to_unicode(self.file_name))
    self.file_txt.SetLabel(to_unicode(self.file_name))
    if (self.info_panel is not None):
      self.panel_sizer.Detach(self.info_panel)
      self.info_panel.Destroy()
    self.info_panel = wx.Panel(self.panel, -1)
    self.panel_sizer.Add(self.info_panel, 1, wx.ALL|wx.EXPAND)
    szr = wx.BoxSizer(wx.VERTICAL)
    self.info_panel.SetSizer(szr)
    box = wx.StaticBox(self.info_panel, -1, "File contents:")
    box_szr = wx.StaticBoxSizer(box, wx.VERTICAL)
    szr.Add(box_szr, 1, wx.EXPAND|wx.ALL, 5)
    grid = wx.FlexGridSizer(cols=2)
    box_szr.Add(grid, 1, wx.EXPAND)
    for line in out.getvalue().splitlines()[1:] :
      label, value = line.strip().split(":")
      txt1 = wx.StaticText(self.info_panel, -1, label + ":")
      font = txt1.GetFont()
      font.SetWeight(wx.FONTWEIGHT_BOLD)
      txt1.SetFont(font)
      grid.Add(txt1, 0, ALN_FLAGS, 5)
      txt2 = wx.StaticText(self.info_panel, -1, value)
      font2 = txt2.GetFont()
      font2.SetFamily(wx.FONTFAMILY_MODERN)
      txt2.SetFont(font2)
      grid.Add(txt2, 0, ALN_FLAGS, 5)
    szr.Fit(self.info_panel)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

class MapCoeffsInfo(InfoPanelBase):
  def __init__(self, *args, **kwds):
    super(MapCoeffsInfo, self).__init__(*args, **kwds)
    self._current_array = None
    self._current_map = None
    box = wx.BoxSizer(wx.HORIZONTAL)
    self.panel_sizer.Add(box, 0)
    bmp = wx.StaticBitmap(self.panel, -1, wxtbx.icons.hkl_file.GetBitmap())
    box.Add(bmp, 0, ALN_FLAGS, 5)
    grid = wx.FlexGridSizer(cols=2)
    box.Add(grid, 0, ALN_FLAGS, 5)
    txt1 = wx.StaticText(self.panel, -1, "File name:")
    font = txt1.GetFont()
    font.SetWeight(wx.FONTWEIGHT_BOLD)
    txt1.SetFont(font)
    grid.Add(txt1, 0, ALN_FLAGS, 5)
    self.file_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.file_txt, 0, ALN_FLAGS, 5)
    txt3 = wx.StaticText(self.panel, -1, "Map coefficients:")
    txt3.SetFont(font)
    grid.Add(txt3, 0, ALN_FLAGS, 5)
    self.map_txt = wx.StaticText(self.panel, -1, "(None)")
    grid.Add(self.map_txt, 0, ALN_FLAGS, 5)

  def set_file(self, file_name):
    self.file_txt.SetLabel(to_unicode(file_name))

  def set_map_coeffs(self, array):
    assert (array.is_complex_array())
    self.map_txt.SetLabel(str(array.info().label_string()))
    self._current_array = array
    map = array.fft_map(resolution_factor=0.25)
    self._current_map = map
    self.info_panel = wx.Panel(self.panel, -1)
    self.panel_sizer.Add(self.info_panel, 1, wx.ALL|wx.EXPAND)
    szr = wx.BoxSizer(wx.VERTICAL)
    self.info_panel.SetSizer(szr)
    box = wx.StaticBox(self.info_panel, -1, "Map info:")
    box_szr = wx.StaticBoxSizer(box, wx.VERTICAL)
    szr.Add(box_szr, 1, wx.EXPAND|wx.ALL, 5)
    grid = wx.FlexGridSizer(cols=2)
    box_szr.Add(grid, 0, wx.EXPAND)
    import iotbx.map_tools
    info_list = iotbx.map_tools.get_map_summary(map)
    for label, value in info_list :
      txt1 = wx.StaticText(self.info_panel, -1, label + ":")
      font = txt1.GetFont()
      font.SetWeight(wx.FONTWEIGHT_BOLD)
      txt1.SetFont(font)
      grid.Add(txt1, 0, ALN_FLAGS, 5)
      str_value = to_str(value)
      alert = False
      if (str_value.endswith("***")):
        str_value = re.sub("\s*\*\*\*", "", str_value)
        alert = True
      txt2 = wx.StaticText(self.info_panel, -1, str_value)
      font2 = txt2.GetFont()
      font2.SetFamily(wx.FONTFAMILY_MODERN)
      if (alert):
        font2.SetWeight(wx.FONTWEIGHT_BOLD)
        txt2.SetForegroundColour((200,0,0))
        txt1.SetForegroundColour((200,0,0))
      txt2.SetFont(font2)
      grid.Add(txt2, 0, ALN_FLAGS, 5)
    self.histogram = wxtbx.plots.histogram(self.info_panel, figure_size=(6,4))
    self.draw_plot(n_bins=20, log_scale=False)
    box_szr.Add(self.histogram, 1, wx.EXPAND, 5)
    hbox = wx.BoxSizer(wx.HORIZONTAL)
    box_szr.Add(hbox, 0, wx.TOP, 5)
    ip = self.info_panel
    btn = wx.Button(ip, -1, "Save plot to file...")
    self.Bind(wx.EVT_BUTTON, lambda evt: self.histogram.save_image(), btn)
    hbox.Add(btn, 0, ALN_FLAGS, 5)
    hbox.Add(wx.StaticText(ip, label="Number of bins:"), 0, ALN_FLAGS, 5)
    self.n_bins_ctrl = wx.SpinCtrl(ip, value="20", min=10, max=50, initial=20)
    self.Bind(wx.EVT_SPINCTRL, self.OnRedraw, self.n_bins_ctrl)
    hbox.Add(self.n_bins_ctrl, 0, ALN_FLAGS, 5)
    self.log_scale_box = wx.CheckBox(ip, label="Use log scale for Y-axis")
    hbox.Add(self.log_scale_box, 0, ALN_FLAGS, 5)
    self.Bind(wx.EVT_CHECKBOX, self.OnRedraw, self.log_scale_box)
    szr.Fit(self.info_panel)
    self.panel.Layout()
    self.panel_sizer.Fit(self.panel)
    self.Fit()

  def OnRedraw(self, event):
    self.draw_plot(
      n_bins = self.n_bins_ctrl.GetValue(),
      log_scale = self.log_scale_box.GetValue())

  def draw_plot(self, n_bins, log_scale):
    self.histogram.figure.clear()
    self.histogram.show_histogram(
      data=self._current_map.real_map(False).as_1d().as_numpy_array(),
      n_bins=n_bins,#n_bins,
      reference_value=0.0,
      x_label="Sigma level",
      y_label="Number of points",
      title="Distribution of grid values (spacing=d_min/4)",
      log_scale=log_scale)
    self.Refresh()

if (__name__ == "__main__"):
  import libtbx.load_env
  pdb1 = libtbx.env.find_in_repositories(
    relative_path="phenix_regression/pdb/2C30.pdb",
    test=os.path.isfile)
  hkl1 = libtbx.env.find_in_repositories(
    relative_path="phenix_regression/wizards/partial_refine_001_map_coeffs.mtz",
    test=os.path.isfile)
  img1 = "/net/cci/dials/jcsg/1vph/data/jcsg/ssrl1/9_2/20040718/TB0723W/11318/11318_2_017.img"
  assert (not None in [pdb1, hkl1])
  app = wx.App(0)
  frame1 = ReflectionFileInfo(None)
  frame1.set_file(hkl1)
  frame1.Show()
  frame2 = PDBFileInfo(None)
  frame2.set_file(pdb1)
  frame2.Show()
  frame2b = PDBChainBisoPanel(None)
  frame2b.set_file(pdb1)
  frame2b.Show()
  if (os.path.exists(img1)):
    frame3 = ImageFileInfo(None)
    frame3.set_file(img1)
    frame3.Show()
  app.MainLoop()
