import numpy as np
import pandas as pd
import enum
import itertools

class Table(object):
  def __init__(self, n_cols):
    self.n_cols = n_cols
    self.rows = []
    self.aligns = ['r' for c in range(n_cols)]

  def get_cell_align(self, r, c):
    align = self.rows[r].cells[c].align
    if align is None:
      return self.aligns[c]
    else:
      return align

  def add_row(self, row):
    if row.ncols() != self.n_cols:
      raise Exception(f'row has invalid number of cols, {row.ncols()} vs. {self.n_cols}')
    self.rows.append(row)

  def empty_row(self):
    return Row.Empty(self.n_cols)

  def expand_rows(self, n_add_cols=1):
    if n_add_cols < 0: raise Exception('n_add_cols has to be positive')
    self.n_cols += n_add_cols
    for row in self.rows:
      row.cells.extend([Cell() for cidx in range(n_add_cols)])

  def add_block(self, data, row=-1, col=0, fmt=None, expand=False):
    if row < 0: row = len(self.rows)
    while len(self.rows) < row + len(data):
      self.add_row(self.empty_row())
    for r in range(len(data)):
      cols = data[r]
      if col + len(cols) > self.n_cols:
        if expand:
          self.expand_rows(col + len(cols) - self.n_cols)
        else:
          raise Exception('number of cols does not fit in table')
      for c in range(len(cols)):
        self.rows[row+r].cells[col+c] = Cell(data[r][c], fmt)

class Row(object):
  def __init__(self, cells, pre_separator=None, post_separator=None):
    self.cells = cells
    self.pre_separator = pre_separator
    self.post_separator = post_separator

  @classmethod
  def Empty(cls, n_cols):
    return Row([Cell() for c in range(n_cols)])

  def add_cell(self, cell):
    self.cells.append(cell)

  def ncols(self):
    return sum([c.span for c in self.cells])



class Color(object):
  def __init__(self, color=(0,0,0), fmt='rgb'):
    if fmt == 'rgb':
      self.color = color
    elif fmt == 'RGB':
      self.color = tuple(c / 255 for c in color)
    else:
      return Exception('invalid color format')

  def as_rgb(self):
    return self.color

  def as_RGB(self):
    return tuple(int(c * 255) for c in self.color)

  @classmethod
  def rgb(cls, r, g, b):
    return Color(color=(r,g,b), fmt='rgb')

  @classmethod
  def RGB(cls, r, g, b):
    return Color(color=(r,g,b), fmt='RGB')


class CellFormat(object):
  def __init__(self, fmt='%s', fgcolor=None, bgcolor=None, bold=False):
    self.fmt = fmt
    self.fgcolor = fgcolor
    self.bgcolor = bgcolor
    self.bold = bold

class Cell(object):
  def __init__(self, data=None, fmt=None, span=1, align=None):
    self.data = data
    if fmt is None:
      fmt = CellFormat()
    self.fmt = fmt
    self.span = span
    self.align = align

  def __str__(self):
    return self.fmt.fmt % self.data

class Separator(enum.Enum):
  HEAD = 1
  BOTTOM = 2
  INNER = 3


class Renderer(object):
  def __init__(self):
    pass

  def cell_str_len(self, cell):
    return len(str(cell))

  def col_widths(self, table):
    widths = [0 for c in range(table.n_cols)]
    for row in table.rows:
      cidx = 0
      for cell in row.cells:
        if cell.span == 1:
          strlen = self.cell_str_len(cell)
          widths[cidx] = max(widths[cidx], strlen)
        cidx += cell.span
    return widths

  def render(self, table):
    raise NotImplementedError('not implemented')

  def __call__(self, table):
    return self.render(table)

  def render_to_file_comment(self):
    return ''

  def render_to_file(self, path, table):
    txt = self.render(table)
    with open(path, 'w') as fp:
      fp.write(txt)

class TerminalRenderer(Renderer):
  def __init__(self, col_sep=' '):
    super().__init__()
    self.col_sep = col_sep

  def render_cell(self, table, row, col, widths):
    cell = table.rows[row].cells[col]
    str = cell.fmt.fmt % cell.data
    str_width = len(str)
    cell_width = sum([widths[idx] for idx in range(col, col+cell.span)])
    cell_width += len(self.col_sep) * (cell.span - 1)
    if len(str) > cell_width:
      str = str[:cell_width]
    if cell.fmt.bold:
      # str = sty.ef.bold + str + sty.rs.bold_dim
      # str = sty.ef.bold + str + sty.rs.bold
      pass
    if cell.fmt.fgcolor is not None:
      # color = cell.fmt.fgcolor.as_RGB()
      # str = sty.fg(*color) + str + sty.rs.fg
      pass
    if str_width < cell_width:
      n_ws = (cell_width - str_width)
      if table.get_cell_align(row, col) == 'r':
        str = ' '*n_ws + str
      elif table.get_cell_align(row, col) == 'l':
        str = str + ' '*n_ws
      elif table.get_cell_align(row, col) == 'c':
        n_ws1 = n_ws // 2
        n_ws0 = n_ws - n_ws1
        str = ' '*n_ws0 + str + ' '*n_ws1
    if cell.fmt.bgcolor is not None:
      # color = cell.fmt.bgcolor.as_RGB()
      # str = sty.bg(*color) + str + sty.rs.bg
      pass
    return str

  def render_separator(self, separator, tab, col_widths, total_width):
    if separator == Separator.HEAD:
      return '='*total_width
    elif separator == Separator.INNER:
      return '-'*total_width
    elif separator == Separator.BOTTOM:
      return '='*total_width

  def render(self, table):
    widths = self.col_widths(table)
    total_width = sum(widths) + len(self.col_sep) * (table.n_cols - 1)
    lines = []
    for ridx, row in enumerate(table.rows):
      if row.pre_separator is not None:
        sepline = self.render_separator(row.pre_separator, table, widths, total_width)
        if len(sepline) > 0:
          lines.append(sepline)
      line = []
      for cidx, cell in enumerate(row.cells):
        line.append(self.render_cell(table, ridx, cidx, widths))
      lines.append(self.col_sep.join(line))
      if row.post_separator is not None:
        sepline = self.render_separator(row.post_separator, table, widths, total_width)
        if len(sepline) > 0:
          lines.append(sepline)
    return '\n'.join(lines)

class MarkdownRenderer(TerminalRenderer):
  def __init__(self):
    super().__init__(col_sep='|')
    self.printed_color_warning = False

  def print_color_warning(self):
    if not self.printed_color_warning:
      print('[WARNING] MarkdownRenderer does not support color yet')
      self.printed_color_warning = True

  def cell_str_len(self, cell):
    strlen = len(str(cell))
    if cell.fmt.bold:
      strlen += 4
    strlen = max(5, strlen)
    return strlen

  def render_cell(self, table, row, col, widths):
    cell = table.rows[row].cells[col]
    str = cell.fmt.fmt % cell.data
    if cell.fmt.bold:
      str = f'**{str}**'

    str_width = len(str)
    cell_width = sum([widths[idx] for idx in range(col, col+cell.span)])
    cell_width += len(self.col_sep) * (cell.span - 1)
    if len(str) > cell_width:
      str = str[:cell_width]
    else:
      n_ws = (cell_width - str_width)
      if table.get_cell_align(row, col) == 'r':
        str = ' '*n_ws + str
      elif table.get_cell_align(row, col) == 'l':
        str = str + ' '*n_ws
      elif table.get_cell_align(row, col) == 'c':
        n_ws1 = n_ws // 2
        n_ws0 = n_ws - n_ws1
        str = ' '*n_ws0 + str + ' '*n_ws1

    if col == 0: str = self.col_sep + str
    if col == table.n_cols - 1: str += self.col_sep

    if cell.fmt.fgcolor is not None:
      self.print_color_warning()
    if cell.fmt.bgcolor is not None:
      self.print_color_warning()
    return str

  def render_separator(self, separator, tab, widths, total_width):
    sep = ''
    if separator == Separator.INNER:
      sep = self.col_sep
      for idx, width in enumerate(widths):
        csep = '-' * (width - 2)
        if tab.get_cell_align(1, idx) == 'r':
          csep = '-' + csep + ':'
        elif tab.get_cell_align(1, idx) == 'l':
          csep = ':' + csep + '-'
        elif tab.get_cell_align(1, idx) == 'c':
          csep = ':' + csep + ':'
        sep += csep + self.col_sep
    return sep


class LatexRenderer(Renderer):
  def __init__(self):
    super().__init__()

  def render_cell(self, table, row, col):
    cell = table.rows[row].cells[col]
    str = cell.fmt.fmt % cell.data
    if cell.fmt.bold:
      str = '{\\bf '+ str + '}'
    if cell.fmt.fgcolor is not None:
      color = cell.fmt.fgcolor.as_rgb()
      str = f'{{\\color[rgb]{{{color[0]},{color[1]},{color[2]}}} ' + str + '}'
    if cell.fmt.bgcolor is not None:
      color = cell.fmt.bgcolor.as_rgb()
      str = f'\\cellcolor[rgb]{{{color[0]},{color[1]},{color[2]}}} ' + str
    align = table.get_cell_align(row, col)
    if cell.span != 1 or align != table.aligns[col]:
      str = f'\\multicolumn{{{cell.span}}}{{{align}}}{{{str}}}'
    return str

  def render_separator(self, separator):
    if separator == Separator.HEAD:
      return '\\toprule'
    elif separator == Separator.INNER:
      return '\\midrule'
    elif separator == Separator.BOTTOM:
      return '\\bottomrule'

  def render(self, table):
    lines = ['\\begin{tabular}{' + ''.join(table.aligns) + '}']
    for ridx, row in enumerate(table.rows):
      if row.pre_separator is not None:
        lines.append(self.render_separator(row.pre_separator))
      line = []
      for cidx, cell in enumerate(row.cells):
        line.append(self.render_cell(table, ridx, cidx))
      lines.append(' & '.join(line) + ' \\\\')
      if row.post_separator is not None:
        lines.append(self.render_separator(row.post_separator))
    lines.append('\\end{tabular}')
    return '\n'.join(lines)

class HtmlRenderer(Renderer):
  def __init__(self, html_class='result_table'):
    super().__init__()
    self.html_class = html_class

  def render_cell(self, table, row, col):
    cell = table.rows[row].cells[col]
    str = cell.fmt.fmt % cell.data
    styles = []
    if cell.fmt.bold:
      styles.append('font-weight: bold;')
    if cell.fmt.fgcolor is not None:
      color = cell.fmt.fgcolor.as_RGB()
      styles.append(f'color: rgb({color[0]},{color[1]},{color[2]});')
    if cell.fmt.bgcolor is not None:
      color = cell.fmt.bgcolor.as_RGB()
      styles.append(f'background-color: rgb({color[0]},{color[1]},{color[2]});')
    align = table.get_cell_align(row, col)
    if align == 'l': align = 'left'
    elif align == 'r': align = 'right'
    elif align == 'c': align = 'center'
    else: raise Exception('invalid align')
    styles.append(f'text-align: {align};')
    row = table.rows[row]
    if row.pre_separator is not None:
      styles.append(f'border-top: {self.render_separator(row.pre_separator)};')
    if row.post_separator is not None:
        styles.append(f'border-bottom: {self.render_separator(row.post_separator)};')
    style = ' '.join(styles)
    str = f'    <td style="{style}" colspan="{cell.span}">{str}</td>\n'
    return str

  def render_separator(self, separator):
    if separator == Separator.HEAD:
      return '1.5pt solid black'
    elif separator == Separator.INNER:
      return '0.75pt solid black'
    elif separator == Separator.BOTTOM:
      return '1.5pt solid black'

  def render(self, table):
    lines = [f'<table width="100%" style="border-collapse: collapse" class={self.html_class}>']
    for ridx, row in enumerate(table.rows):
      line = [f'  <tr>\n']
      for cidx, cell in enumerate(row.cells):
        line.append(self.render_cell(table, ridx, cidx))
      line.append('  </tr>\n')
      lines.append(' '.join(line))
    lines.append('</table>')
    return '\n'.join(lines)


def pandas_to_table(rowname, colname, valname, data, val_cell_fmt=CellFormat(fmt='%.4f'), best_val_cell_fmt=CellFormat(fmt='%.4f', bold=True), best_is_max=[]):
  rnames = data[rowname].unique()
  cnames = data[colname].unique()
  tab = Table(1+len(cnames))

  header = [Cell('', align='r')]
  header.extend([Cell(h, align='r') for h in cnames])
  header = Row(header, pre_separator=Separator.HEAD, post_separator=Separator.INNER)
  tab.add_row(header)

  for rname in rnames:
    cells = [Cell(rname, align='l')]
    for cname in cnames:
      cdata = data[data[colname] == cname]
      if cname in best_is_max:
        bestval = cdata[valname].max()
        val = cdata[cdata[rowname] == rname][valname].max()
      else:
        bestval = cdata[valname].min()
        val = cdata[cdata[rowname] == rname][valname].min()
      if val == bestval:
        fmt = best_val_cell_fmt
      else:
        fmt = val_cell_fmt
      cells.append(Cell(val, align='r', fmt=fmt))
    tab.add_row(Row(cells))
  tab.rows[-1].post_separator = Separator.BOTTOM
  return tab



if __name__ == '__main__':
  # df = pd.read_pickle('full.df')
  # best_is_max = ['movF0.5', 'movF1.0']
  # tab = pandas_to_table(rowname='method', colname='metric', valname='val', data=df, best_is_max=best_is_max)

  # renderer = TerminalRenderer()
  # print(renderer(tab))

  tab = Table(7)
  # header = Row([Cell('header', span=7, align='c')], pre_separator=Separator.HEAD, post_separator=Separator.INNER)
  # tab.add_row(header)
  # header2 = Row([Cell('thisisaverylongheader', span=4, align='c'), Cell('vals2', span=3, align='c')], post_separator=Separator.INNER)
  # tab.add_row(header2)
  tab.add_row(Row([Cell(f'c{c}') for c in range(7)]))
  tab.rows[-1].post_separator = Separator.INNER
  tab.add_block(np.arange(15*7).reshape(15,7))
  tab.rows[4].cells[2].fmt = CellFormat(bold=True)
  tab.rows[2].cells[1].fmt = CellFormat(fgcolor=Color.rgb(0.2,0.6,0.1))
  tab.rows[2].cells[2].fmt = CellFormat(bgcolor=Color.rgb(0.7,0.1,0.5))
  tab.rows[5].cells[3].fmt = CellFormat(bold=True,bgcolor=Color.rgb(0.7,0.1,0.5),fgcolor=Color.rgb(0.1,0.1,0.1))
  tab.rows[-1].post_separator = Separator.BOTTOM

  renderer = TerminalRenderer()
  print(renderer(tab))
  renderer = MarkdownRenderer()
  print(renderer(tab))

  # renderer = HtmlRenderer()
  # html_tab = renderer(tab)
  # print(html_tab)
  # with open('test.html', 'w') as fp:
  #   fp.write(html_tab)

  # import latex

  # renderer = LatexRenderer()
  # ltx_tab = renderer(tab)
  # print(ltx_tab)

  # with open('test.tex', 'w') as fp:
  #   latex.write_doc_prefix(fp, document_class='article')
  #   fp.write('this is text that should appear before the table and should be long enough to wrap around.\n'*40)
  #   fp.write('\\begin{table}')
  #   fp.write(ltx_tab)
  #   fp.write('\\end{table}')
  #   fp.write('this is text that should appear after the table and should be long enough to wrap around.\n'*40)
  #   latex.write_doc_suffix(fp)