"""
Provides a simple table class.  A SimpleTable is essentially
a list of lists plus some formatting functionality.

Dependencies: the Python 2.5+ standard library.

Installation: just copy this module into your working directory (or
   anywhere in your pythonpath).

Basic use::

   mydata = [[11,12],[21,22]]  # data MUST be 2-dimensional
   myheaders = [ "Column 1", "Column 2" ]
   mystubs = [ "Row 1", "Row 2" ]
   tbl = SimpleTable(mydata, myheaders, mystubs, title="Title")
   print( tbl )
   print( tbl.as_csv() )

A SimpleTable is inherently (but not rigidly) rectangular.
You should create it from a *rectangular* (2d!) iterable of data.
Each item in your rectangular iterable will become the data
of a single Cell.  In principle, items can be any object,
not just numbers and strings.  However, default conversion
during table production is by simple string interpolation.
(So you cannot have a tuple as a data item *and* rely on
the default conversion.)

A SimpleTable allows only one column (the first) of stubs at
initilization, concatenation of tables allows you to produce tables
with interior stubs.  (You can also assign the datatype 'stub' to the
cells in any column, or use ``insert_stubs``.) A SimpleTable can be
concatenated with another SimpleTable or extended by another
SimpleTable. ::

    table1.extend_right(table2)
    table1.extend(table2)

A SimpleTable can be initialized with `datatypes`: a list of ints that
provide indexes into `data_fmts` and `data_aligns`.  Each data cell is
assigned a datatype, which will control formatting.  If you do not
specify the `datatypes` list, it will be set to ``range(ncols)`` where
`ncols` is the number of columns in the data.  (I.e., cells in a
column have their own datatype.) This means that you can just specify
`data_fmts` without bothering to provide a `datatypes` list.  If
``len(datatypes)<ncols`` then datatype assignment will cycle across a
row.  E.g., if you provide 10 columns of data with ``datatypes=[0,1]``
then you will have 5 columns of datatype 0 and 5 columns of datatype
1, alternating.  Corresponding to this specification, you should provide
a list of two ``data_fmts`` and a list of two ``data_aligns``.

Cells can be assigned labels as their `datatype` attribute.
You can then provide a format for that lable.
Us the SimpleTable's `label_cells` method to do this.  ::

    def mylabeller(cell):
        if cell.data is np.nan:
            return 'missing'

    mytable.label_cells(mylabeller)
    print(mytable.as_text(missing='-'))

Potential problems for Python 3
-------------------------------

- Calls ``next`` instead of ``__next__``.
  The 2to3 tool should handle that no problem.
  (We will switch to the `next` function if 2.5 support is ever dropped.)
- Let me know if you find other problems.

:contact: alan dot isaac at gmail dot com
:requires: Python 2.5.1+
:note: current version
:note: HTML data format currently specifies tags
:todo: support a bit more of http://www.oasis-open.org/specs/tr9503.html
:todo: add labels2formatters method, that associates a cell formatter with a
       datatype
:todo: add colspan support to Cell
:since: 2008-12-21
:change: 2010-05-02 eliminate newlines that came before and after table
:change: 2010-05-06 add `label_cells` to `SimpleTable`
"""

from statsmodels.compat.python import lmap, lrange

from itertools import cycle, zip_longest
import csv

def csv2st(csvfile, headers=False, stubs=False, title=None):
    """Return SimpleTable instance,
    created from the data in `csvfile`,
    which is in comma separated values format.
    The first row may contain headers: set headers=True.
    The first column may contain stubs: set stubs=True.
    Can also supply headers and stubs as tuples of strings.
    """
    rows = list()
    with open(csvfile, 'r', encoding="utf-8") as fh:
        reader = csv.reader(fh)
        if headers is True:
            headers = next(reader)
        elif headers is False:
            headers = ()
        if stubs is True:
            stubs = list()
            for row in reader:
                if row:
                    stubs.append(row[0])
                    rows.append(row[1:])
        else:  # no stubs, or stubs provided
            for row in reader:
                if row:
                    rows.append(row)
        if stubs is False:
            stubs = ()
    ncols = len(rows[0])
    if any(len(row) != ncols for row in rows):
        raise IOError('All rows of CSV file must have same length.')
    return SimpleTable(data=rows, headers=headers, stubs=stubs)

class SimpleTable(list):
    """Produce a simple ASCII, CSV, HTML, or LaTeX table from a
    *rectangular* (2d!) array of data, not necessarily numerical.
    Directly supports at most one header row,
    which should be the length of data[0].
    Directly supports at most one stubs column,
    which must be the length of data.
    (But see `insert_stubs` method.)
    See globals `default_txt_fmt`, `default_csv_fmt`, `default_html_fmt`,
    and `default_latex_fmt` for formatting options.

    Sample uses::

        mydata = [[11,12],[21,22]]  # data MUST be 2-dimensional
        myheaders = [ "Column 1", "Column 2" ]
        mystubs = [ "Row 1", "Row 2" ]
        tbl = text.SimpleTable(mydata, myheaders, mystubs, title="Title")
        print( tbl )
        print( tbl.as_html() )
        # set column specific data formatting
        tbl = text.SimpleTable(mydata, myheaders, mystubs,
            data_fmts=["%3.2f","%d"])
        print( tbl.as_csv() )
        with open('c:/temp/temp.tex','w') as fh:
            fh.write( tbl.as_latex_tabular() )
    """
    def __init__(self, data, headers=None, stubs=None, title='',
                 datatypes=None, csv_fmt=None, txt_fmt=None, ltx_fmt=None,
                 html_fmt=None, celltype=None, rowtype=None, **fmt_dict):
        """
        Parameters
        ----------
        data : list of lists or 2d array (not matrix!)
            R rows by K columns of table elements
        headers : list (or tuple) of str
            sequence of K strings, one per header
        stubs : list (or tuple) of str
            sequence of R strings, one per stub
        title : str
            title of the table
        datatypes : list of int
            indexes to `data_fmts`
        txt_fmt : dict
            text formatting options
        ltx_fmt : dict
            latex formatting options
        csv_fmt : dict
            csv formatting options
        hmtl_fmt : dict
            hmtl formatting options
        celltype : class
            the cell class for the table (default: Cell)
        rowtype : class
            the row class for the table (default: Row)
        fmt_dict : dict
            general formatting options
        """
        self.title = title
        self._datatypes = datatypes
        if self._datatypes is None:
            self._datatypes = [] if len(data) == 0 else lrange(len(data[0]))
        # start with default formatting
        self._txt_fmt = default_txt_fmt.copy()
        self._latex_fmt = default_latex_fmt.copy()
        self._csv_fmt = default_csv_fmt.copy()
        self._html_fmt = default_html_fmt.copy()
        # substitute any general user specified formatting
        # :note: these will be overridden by output specific arguments
        self._csv_fmt.update(fmt_dict)
        self._txt_fmt.update(fmt_dict)
        self._latex_fmt.update(fmt_dict)
        self._html_fmt.update(fmt_dict)
        # substitute any output-type specific formatting
        self._csv_fmt.update(csv_fmt or dict())
        self._txt_fmt.update(txt_fmt or dict())
        self._latex_fmt.update(ltx_fmt or dict())
        self._html_fmt.update(html_fmt or dict())
        self.output_formats = dict(
            txt=self._txt_fmt,
            csv=self._csv_fmt,
            html=self._html_fmt,
            latex=self._latex_fmt
        )
        self._Cell = celltype or Cell
        self._Row = rowtype or Row
        rows = self._data2rows(data)  # a list of Row instances
        list.__init__(self, rows)
        self._add_headers_stubs(headers, stubs)
        self._colwidths = dict()

    def __str__(self):
        return self.as_text()

    def __repr__(self):
        return str(type(self))

    def _repr_html_(self, **fmt_dict):
        return self.as_html(**fmt_dict)

    def _repr_latex_(self, center=True, **fmt_dict):
        return self.as_latex_tabular(center, **fmt_dict)

    def _add_headers_stubs(self, headers, stubs):
        """Return None.  Adds headers and stubs to table,
        if these were provided at initialization.
        Parameters
        ----------
        headers : list[str]
            K strings, where K is number of columns
        stubs : list[str]
            R strings, where R is number of non-header rows

        :note: a header row does not receive a stub!
        """
        if headers:
            self.insert_header_row(0, headers, dec_below='header_dec_below')
        if stubs:
            self.insert_stubs(0, stubs)

    def insert(self, idx, row, datatype=None):
        """Return None.  Insert a row into a table.
        """
        if datatype is None:
            try:
                datatype = row.datatype
            except AttributeError:
                pass
        row = self._Row(row, datatype=datatype, table=self)
        list.insert(self, idx, row)

    def insert_header_row(self, rownum, headers, dec_below='header_dec_below'):
        """Return None.  Insert a row of headers,
        where ``headers`` is a sequence of strings.
        (The strings may contain newlines, to indicated multiline headers.)
        """
        header_rows = [header.split('\n') for header in headers]
        # rows in reverse order
        rows = list(zip_longest(*header_rows, **dict(fillvalue='')))
        rows.reverse()
        for i, row in enumerate(rows):
            self.insert(rownum, row, datatype='header')
            if i == 0:
                self[rownum].dec_below = dec_below
            else:
                self[rownum].dec_below = None

    def insert_stubs(self, loc, stubs):
        """Return None.  Insert column of stubs at column `loc`.
        If there is a header row, it gets an empty cell.
        So ``len(stubs)`` should equal the number of non-header rows.
        """
        _Cell = self._Cell
        stubs = iter(stubs)
        for row in self:
            if row.datatype == 'header':
                empty_cell = _Cell('', datatype='empty')
                row.insert(loc, empty_cell)
            else:
                try:
                    row.insert_stub(loc, next(stubs))
                except StopIteration:
                    raise ValueError('length of stubs must match table length')

    def _data2rows(self, raw_data):
        """Return list of Row,
        the raw data as rows of cells.
        """

        _Cell = self._Cell
        _Row = self._Row
        rows = []
        for datarow in raw_data:
            dtypes = cycle(self._datatypes)
            newrow = _Row(datarow, datatype='data', table=self, celltype=_Cell)
            for cell in newrow:
                cell.datatype = next(dtypes)
                cell.row = newrow  # a cell knows its row
            rows.append(newrow)

        return rows

    def pad(self, s, width, align):
        """DEPRECATED: just use the pad function"""
        return pad(s, width, align)

    def _get_colwidths(self, output_format, **fmt_dict):
        """Return list, the calculated widths of each column."""
        output_format = get_output_format(output_format)
        fmt = self.output_formats[output_format].copy()
        fmt.update(fmt_dict)
        ncols = max(len(row) for row in self)
        request = fmt.get('colwidths')
        if request == 0:  # assume no extra space desired (e.g, CSV)
            return [0] * ncols
        elif request is None:  # assume no extra space desired (e.g, CSV)
            request = [0] * ncols
        elif isinstance(request, int):
            request = [request] * ncols
        elif len(request) < ncols:
            request = [request[i % len(request)] for i in range(ncols)]
        min_widths = []
        for col in zip(*self):
            maxwidth = max(len(c.format(0, output_format, **fmt)) for c in col)
            min_widths.append(maxwidth)
        result = lmap(max, min_widths, request)
        return result

    def get_colwidths(self, output_format, **fmt_dict):
        """Return list, the widths of each column."""
        call_args = [output_format]
        for k, v in sorted(fmt_dict.items()):
            if isinstance(v, list):
                call_args.append((k, tuple(v)))
            elif isinstance(v, dict):
                call_args.append((k, tuple(sorted(v.items()))))
            else:
                call_args.append((k, v))
        key = tuple(call_args)
        try:
            return self._colwidths[key]
        except KeyError:
            self._colwidths[key] = self._get_colwidths(output_format,
                                                       **fmt_dict)
            return self._colwidths[key]

    def _get_fmt(self, output_format, **fmt_dict):
        """Return dict, the formatting options.
        """
        output_format = get_output_format(output_format)
        # first get the default formatting
        try:
            fmt = self.output_formats[output_format].copy()
        except KeyError:
            raise ValueError('Unknown format: %s' % output_format)
        # then, add formatting specific to this call
        fmt.update(fmt_dict)
        return fmt

    def as_csv(self, **fmt_dict):
        """Return string, the table in CSV format.
        Currently only supports comma separator."""
        # fetch the format, which may just be default_csv_format
        fmt = self._get_fmt('csv', **fmt_dict)
        return self.as_text(**fmt)

    def as_text(self, **fmt_dict):
        """Return string, the table as text."""
        # fetch the text format, override with fmt_dict
        fmt = self._get_fmt('txt', **fmt_dict)
        # get rows formatted as strings
        formatted_rows = [row.as_string('text', **fmt) for row in self]
        rowlen = len(formatted_rows[-1])  # do not use header row

        # place decoration above the table body, if desired
        table_dec_above = fmt.get('table_dec_above', '=')
        if table_dec_above:
            formatted_rows.insert(0, table_dec_above * rowlen)
        # next place a title at the very top, if desired
        # :note: user can include a newlines at end of title if desired
        title = self.title
        if title:
            title = pad(self.title, rowlen, fmt.get('title_align', 'c'))
            formatted_rows.insert(0, title)
        # add decoration below the table, if desired
        table_dec_below = fmt.get('table_dec_below', '-')
        if table_dec_below:
            formatted_rows.append(table_dec_below * rowlen)
        return '\n'.join(formatted_rows)

    def as_html(self, **fmt_dict):
        """Return string.
        This is the default formatter for HTML tables.
        An HTML table formatter must accept as arguments
        a table and a format dictionary.
        """
        # fetch the text format, override with fmt_dict
        fmt = self._get_fmt('html', **fmt_dict)
        formatted_rows = ['
']
        if self.title:
            title = '
%s
' % self.title
            formatted_rows.append(title)
        formatted_rows.extend(row.as_string('html', **fmt) for row in self)
        formatted_rows.append('
')
        return '\n'.join(formatted_rows)

    def as_latex_tabular(self, center=True, **fmt_dict):
        '''Return string, the table as a LaTeX tabular environment.
        Note: will require the booktabs package.'''
        # fetch the text format, override with fmt_dict
        fmt = self._get_fmt('latex', **fmt_dict)

        formatted_rows = []
        if center:
            formatted_rows.append(r'\begin{center}')

        table_dec_above = fmt['table_dec_above'] or ''
        table_dec_below = fmt['table_dec_below'] or ''

        prev_aligns = None
        last = None
        for row in self + [last]:
            if row == last:
                aligns = None
            else:
                aligns = row.get_aligns('latex', **fmt)

            if aligns != prev_aligns:
                # When the number/type of columns changes...
                if prev_aligns:
                    # ... if there is a tabular to close, close it...
                    formatted_rows.append(table_dec_below)
                    formatted_rows.append(r'\end{tabular}')
                if aligns:
                    # ... and if there are more lines, open a new one:
                    formatted_rows.append(r'\begin{tabular}{%s}' % aligns)
                    if not prev_aligns:
                        # (with a nice line if it's the top of the whole table)
                        formatted_rows.append(table_dec_above)
            if row != last:
                formatted_rows.append(
                    row.as_string(output_format='latex', **fmt))
            prev_aligns = aligns
        # tabular does not support caption, but make it available for
        # figure environment
        if self.title:
            title = r'%%\caption{%s}' % self.title
            formatted_rows.append(title)
        if center:
            formatted_rows.append(r'\end{center}')

        # Replace $$ due to bug in GH 5444
        return '\n'.join(formatted_rows).replace('$$', ' ')

    def extend_right(self, table):
        """Return None.
        Extend each row of `self` with corresponding row of `table`.
        Does **not** import formatting from ``table``.
        This generally makes sense only if the two tables have
        the same number of rows, but that is not enforced.
        :note: To extend append a table below, just use `extend`,
        which is the ordinary list method.  This generally makes sense
        only if the two tables have the same number of columns,
        but that is not enforced.
        """
        for row1, row2 in zip(self, table):
            row1.extend(row2)

    def label_cells(self, func):
        """Return None.  Labels cells based on `func`.
        If ``func(cell) is None`` then its datatype is
        not changed; otherwise it is set to ``func(cell)``.
        """
        for row in self:
            for cell in row:
                label = func(cell)
                if label is not None:
                    cell.datatype = label

    @property
    def data(self):
        return [row.data for row in self]

def pad(s, width, align):
    """Return string padded with spaces,
    based on alignment parameter."""
    if align == 'l':
        s = s.ljust(width)
    elif align == 'r':
        s = s.rjust(width)
    else:
        s = s.center(width)
    return s

class Row(list):
    """Provides a table row as a list of cells.
    A row can belong to a SimpleTable, but does not have to.
    """
    def __init__(self, seq, datatype='data', table=None, celltype=None,
                 dec_below='row_dec_below', **fmt_dict):
        """
        Parameters
        ----------
        seq : sequence of data or cells
        table : SimpleTable
        datatype : str ('data' or 'header')
        dec_below : str
          (e.g., 'header_dec_below' or 'row_dec_below')
          decoration tag, identifies the decoration to go below the row.
          (Decoration is repeated as needed for text formats.)
        """
        self.datatype = datatype
        self.table = table
        if celltype is None:
            if table is None:
                celltype = Cell
            else:
                celltype = table._Cell
        self._Cell = celltype
        self._fmt = fmt_dict
        self.special_fmts = dict()  # special formatting for any output format
        self.dec_below = dec_below
        list.__init__(self, (celltype(cell, row=self) for cell in seq))

    def add_format(self, output_format, **fmt_dict):
        """
        Return None. Adds row-instance specific formatting
        for the specified output format.
        Example: myrow.add_format('txt', row_dec_below='+-')
        """
        output_format = get_output_format(output_format)
        if output_format not in self.special_fmts:
            self.special_fmts[output_format] = dict()
        self.special_fmts[output_format].update(fmt_dict)

    def insert_stub(self, loc, stub):
        """Return None.  Inserts a stub cell
        in the row at `loc`.
        """
        _Cell = self._Cell
        if not isinstance(stub, _Cell):
            stub = stub
            stub = _Cell(stub, datatype='stub', row=self)
        self.insert(loc, stub)

    def _get_fmt(self, output_format, **fmt_dict):
        """Return dict, the formatting options.
        """
        output_format = get_output_format(output_format)
        # first get the default formatting
        try:
            fmt = default_fmts[output_format].copy()
        except KeyError:
            raise ValueError('Unknown format: %s' % output_format)
        # second get table specific formatting (if possible)
        try:
            fmt.update(self.table.output_formats[output_format])
        except AttributeError:
            pass
        # finally, add formatting for this row and this call
        fmt.update(self._fmt)
        fmt.update(fmt_dict)
        special_fmt = self.special_fmts.get(output_format, None)
        if special_fmt is not None:
            fmt.update(special_fmt)
        return fmt

    def get_aligns(self, output_format, **fmt_dict):
        """Return string, sequence of column alignments.
        Ensure comformable data_aligns in `fmt_dict`."""
        fmt = self._get_fmt(output_format, **fmt_dict)
        return ''.join(cell.alignment(output_format, **fmt) for cell in self)

    def as_string(self, output_format='txt', **fmt_dict):
        """Return string: the formatted row.
        This is the default formatter for rows.
        Override this to get different formatting.
        A row formatter must accept as arguments
        a row (self) and an output format,
        one of ('html', 'txt', 'csv', 'latex').
        """
        fmt = self._get_fmt(output_format, **fmt_dict)

        # get column widths
        try:
            colwidths = self.table.get_colwidths(output_format, **fmt)
        except AttributeError:
            colwidths = fmt.get('colwidths')
        if colwidths is None:
            colwidths = (0,) * len(self)

        colsep = fmt['colsep']
        row_pre = fmt.get('row_pre', '')
        row_post = fmt.get('row_post', '')
        formatted_cells = []
        for cell, width in zip(self, colwidths):
            content = cell.format(width, output_format=output_format, **fmt)
            formatted_cells.append(content)
        formatted_row = row_pre + colsep.join(formatted_cells) + row_post
        formatted_row = self._decorate_below(formatted_row, output_format,
                                             **fmt)
        return formatted_row

    def _decorate_below(self, row_as_string, output_format, **fmt_dict):
        """This really only makes sense for the text and latex output formats.
        """
        dec_below = fmt_dict.get(self.dec_below, None)
        if dec_below is None:
            result = row_as_string
        else:
            output_format = get_output_format(output_format)
            if output_format == 'txt':
                row0len = len(row_as_string)
                dec_len = len(dec_below)
                repeat, addon = divmod(row0len, dec_len)
                result = row_as_string + "\n" + (dec_below * repeat +
                                                 dec_below[:addon])
            elif output_format == 'latex':
                result = row_as_string + "\n" + dec_below
            else:
                raise ValueError("I cannot decorate a %s header." %
                                 output_format)
        return result

    @property
    def data(self):
        return [cell.data for cell in self]

class Cell:
    """Provides a table cell.
    A cell can belong to a Row, but does not have to.
    """
    def __init__(self, data='', datatype=None, row=None, **fmt_dict):
        if isinstance(data, Cell):
            # might have passed a Cell instance
            self.data = data.data
            self._datatype = data.datatype
            self._fmt = data._fmt
        else:
            self.data = data
            self._datatype = datatype
            self._fmt = dict()
        self._fmt.update(fmt_dict)
        self.row = row

    def __str__(self):
        return '%s' % self.data

    def _get_fmt(self, output_format, **fmt_dict):
        """Return dict, the formatting options.
        """
        output_format = get_output_format(output_format)
        # first get the default formatting
        try:
            fmt = default_fmts[output_format].copy()
        except KeyError:
            raise ValueError('Unknown format: %s' % output_format)
        # then get any table specific formtting
        try:
            fmt.update(self.row.table.output_formats[output_format])
        except AttributeError:
            pass
        # then get any row specific formtting
        try:
            fmt.update(self.row._fmt)
        except AttributeError:
            pass
        # finally add formatting for this instance and call
        fmt.update(self._fmt)
        fmt.update(fmt_dict)
        return fmt

    def alignment(self, output_format, **fmt_dict):
        fmt = self._get_fmt(output_format, **fmt_dict)
        datatype = self.datatype
        data_aligns = fmt.get('data_aligns', 'c')
        if isinstance(datatype, int):
            align = data_aligns[datatype % len(data_aligns)]
        elif datatype == 'stub':
            # still support deprecated `stubs_align`
            align = fmt.get('stubs_align') or fmt.get('stub_align', 'l')
        elif datatype in fmt:
            label_align = '%s_align' % datatype
            align = fmt.get(label_align, 'c')
        else:
            raise ValueError('Unknown cell datatype: %s' % datatype)
        return align

    @staticmethod
    def _latex_escape(data, fmt, output_format):
        if output_format != 'latex':
            return data
        if "replacements" in fmt:
            if isinstance(data, str):
                for repl in sorted(fmt["replacements"]):
                    data = data.replace(repl, fmt["replacements"][repl])
        return data

    def format(self, width, output_format='txt', **fmt_dict):
        """Return string.
        This is the default formatter for cells.
        Override this to get different formating.
        A cell formatter must accept as arguments
        a cell (self) and an output format,
        one of ('html', 'txt', 'csv', 'latex').
        It will generally respond to the datatype,
        one of (int, 'header', 'stub').
        """
        fmt = self._get_fmt(output_format, **fmt_dict)

        data = self.data
        datatype = self.datatype
        data_fmts = fmt.get('data_fmts')
        if data_fmts is None:
            # chk allow for deprecated use of data_fmt
            data_fmt = fmt.get('data_fmt')
            if data_fmt is None:
                data_fmt = '%s'
            data_fmts = [data_fmt]
        if isinstance(datatype, int):
            datatype = datatype % len(data_fmts)  # constrain to indexes
            data_fmt = data_fmts[datatype]
            if isinstance(data_fmt, str):
                content = data_fmt % (data,)
            elif callable(data_fmt):
                content = data_fmt(data)
            else:
                raise TypeError("Must be a string or a callable")
            if datatype == 0:
                content = self._latex_escape(content, fmt, output_format)
        elif datatype in fmt:
            data = self._latex_escape(data, fmt, output_format)

            dfmt = fmt.get(datatype)
            try:
                content = dfmt % (data,)
            except TypeError:  # dfmt is not a substitution string
                content = dfmt
        else:
            raise ValueError('Unknown cell datatype: %s' % datatype)
        align = self.alignment(output_format, **fmt)
        return pad(content, width, align)

    def get_datatype(self):
        if self._datatype is None:
            dtype = self.row.datatype
        else:
            dtype = self._datatype
        return dtype

    def set_datatype(self, val):
        # TODO: add checking
        self._datatype = val
    datatype = property(get_datatype, set_datatype)

# begin: default formats for SimpleTable
""" Some formatting suggestions:

- if you want rows to have no extra spacing,
  set colwidths=0 and colsep=''.
  (Naturally the columns will not align.)
- if you want rows to have minimal extra spacing,
  set colwidths=1.  The columns will align.
- to get consistent formatting, you should leave
  all field width handling to SimpleTable:
  use 0 as the field width in data_fmts.  E.g., ::

        data_fmts = ["%#0.6g","%#0.6g","%#0.4g","%#0.4g"],
        colwidths = 14,
        data_aligns = "r",
"""
default_txt_fmt = dict(
    fmt='txt',
    # basic table formatting
    table_dec_above='=',
    table_dec_below='-',
    title_align='c',
    # basic row formatting
    row_pre='',
    row_post='',
    header_dec_below='-',
    row_dec_below=None,
    colwidths=None,
    colsep=' ',
    data_aligns="r",  # GH 1477
    # data formats
    # data_fmt="%s",  #deprecated; use data_fmts
    data_fmts=["%s"],
    # labeled alignments
    # stubs_align='l',   #deprecated; use data_fmts
    stub_align='l',
    header_align='c',
    # labeled formats
    header_fmt='%s',  # deprecated; just use 'header'
    stub_fmt='%s',  # deprecated; just use 'stub'
    header='%s',
    stub='%s',
    empty_cell='',  # deprecated; just use 'empty'
    empty='',
    missing='--',
)

default_csv_fmt = dict(
    fmt='csv',
    table_dec_above=None,  # '',
    table_dec_below=None,  # '',
    # basic row formatting
    row_pre='',
    row_post='',
    header_dec_below=None,  # '',
    row_dec_below=None,
    title_align='',
    data_aligns="l",
    colwidths=None,
    colsep=',',
    # data formats
    data_fmt='%s',  # deprecated; use data_fmts
    data_fmts=['%s'],
    # labeled alignments
    # stubs_align='l',   # deprecated; use data_fmts
    stub_align="l",
    header_align='c',
    # labeled formats
    header_fmt='"%s"',  # deprecated; just use 'header'
    stub_fmt='"%s"',  # deprecated; just use 'stub'
    empty_cell='',  # deprecated; just use 'empty'
    header='%s',
    stub='%s',
    empty='',
    missing='--',
)

default_html_fmt = dict(
    # basic table formatting
    table_dec_above=None,
    table_dec_below=None,
    header_dec_below=None,
    row_dec_below=None,
    title_align='c',
    # basic row formatting
    colwidths=None,
    colsep=' ',
    row_pre='
\n  ',
    row_post='\n
',
    data_aligns="c",
    # data formats
    data_fmts=['%s	'],
    data_fmt="%s	",  # deprecated; use data_fmts
    # labeled alignments
    # stubs_align='l',   #deprecated; use data_fmts
    stub_align='l',
    header_align='c',
    # labeled formats
    header_fmt='%s	',  # deprecated; just use `header`
    stub_fmt='%s	',  # deprecated; just use `stub`
    empty_cell='	',  # deprecated; just use `empty`
    header='%s	',
    stub='%s	',
    empty='	',
    missing='--	',
)

default_latex_fmt = dict(
    fmt='ltx',
    # basic table formatting
    table_dec_above=r'\toprule',
    table_dec_below=r'\bottomrule',
    header_dec_below=r'\midrule',
    row_dec_below=None,
    strip_backslash=True,  # NotImplemented
    # row formatting
    row_post=r'  \\',
    data_aligns='c',
    colwidths=None,
    colsep=' & ',
    # data formats
    data_fmts=['%s'],
    data_fmt='%s',  # deprecated; use data_fmts
    # labeled alignments
    # stubs_align='l',   # deprecated; use data_fmts
    stub_align='l',
    header_align='c',
    empty_align='l',
    # labeled formats
    header_fmt=r'\textbf{%s}',  # deprecated; just use 'header'
    stub_fmt=r'\textbf{%s}',  # deprecated; just use 'stub'
    empty_cell='',  # deprecated; just use 'empty'
    header=r'\textbf{%s}',
    stub=r'\textbf{%s}',
    empty='',
    missing='--',
    # replacements will be processed in lexicographical order
    replacements={"#": r"\#",
                  "$": r"\$",
                  "%": r"\%",
                  "&": r"\&",
                  ">": r"$>$",
                  "_": r"\_",
                  "|": r"$|$"}
)

default_fmts = dict(
    html=default_html_fmt,
    txt=default_txt_fmt,
    latex=default_latex_fmt,
    csv=default_csv_fmt
)
output_format_translations = dict(
    htm='html',
    text='txt',
    ltx='latex'
)

def get_output_format(output_format):
    if output_format not in ('html', 'txt', 'latex', 'csv'):
        try:
            output_format = output_format_translations[output_format]
        except KeyError:
            raise ValueError('unknown output format %s' % output_format)
    return output_format


from typing import DefaultDict
import numpy as np
from numpy.core.fromnumeric import size
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import pandas as pd
from sklearn.metrics import roc_auc_score,log_loss
import os
import random
from utils import Config
import wandb
import collections
import math
import scipy.stats
from torch.optim import Adagrad

if __name__ == "__main__":

    # preprocess
    print(f'num_feat:{num_categorial_feat}')
    examplar_size = data_num//task_num
    criterion = nn.BCELoss()
    wconfig = {
        "weight_decay":weight_decay,
        "batch_win_size":Config.batch_win_size,
        "lr":lr,
        "pretrain_epoch":epoch_num,
        "finetune_epoch":epoch_num,
        "patience":PATIENCE,
        "task_num":task_num,
        "batch_size":batch_size,
        "hidden_dim":HID_DIM,
        "device":device
    }
    wandb.init(project="old_sampler", entity="ludlethwang",config=wconfig,name=f"icg")

    # training & evaulation for all configuration
    for ifrandom, oldrate in [(0,2)]:
        model = BaseModel(sizes).to(device)
        embeds = [eb.weight.detach() for eb in model.embeddings]
        print(f'model param nums: {sum([param.numel() for param in model.parameters()])}')
        print(f"batch_size:{batch_size}")
        optimizer = Adagrad(
            model.parameters(),
            lr=lr,
            initial_accumulator_value=1e-8,
            weight_decay=weight_decay
        )
        examplar_features, examplar_labels = [f[:int(data_num/task_num*Config.batch_win_size)] for f in categorial_features], labels[:int(data_num/task_num*Config.batch_win_size)]
        guards = [torch.zeros([sizes[_]+1],dtype=torch.long) for _ in range(num_categorial_feat)]
        overall_auc = []
        overall_logloss = []
        overall_auc_old = []
        overall_logloss_old = []
        overall_auc_notold = []
        overall_logloss_notold = []

        # incremental learning train and evaluate $win_size tasks one by one
        for task_idx in range(Config.batch_win_size, task_num-1):
            global_auc = 0.0
            global_logloss = 10.0
            global_auc_old = 0.0
            global_logloss_old = 10.0
            global_auc_notold = 0.0
            global_logloss_notold = 10.0
            global_patience = PATIENCE
            patience = global_patience
            print(f'running task {task_idx}')

            # sampling the examplar training set
            dataloader,test_loader,guards,examplar_features,examplar_labels = olddata_sampler(
                task_idx,ifrandom,oldrate,1,examplar_features,examplar_labels,guards
                )
            _,test_loader_old,__,___,____ = olddata_sampler(task_idx,ifrandom,oldrate,0,None,None,None)
            _,test_loader_notold,__,___,____ = olddata_sampler(task_idx,ifrandom,oldrate,2,None,None,None)
            for epoch in range(epoch_num):
                model.train()
                total_loss = []

                #training process
                for features, batchlabel in dataloader:
                    features = features.to(device)
                    batchlabel = batchlabel.to(device)
                    prediction =model(features)
                    if if_guard:
                        last_embeds = embeds
                        embeds = [eb.weight.detach().clone() for eb in model.embeddings]
                        reg_l = model.reg(guards,last_embeds)
                        # print(reg_l)
                    else:
                        reg_l = 0
                    loss = criterion(prediction.view(-1),batchlabel)+reg_l
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    total_loss.append(loss.item())
                train_loss = sum(total_loss)/len(total_loss)
                if (epoch+1)%2==0:
                    print(epoch,"train loss:",train_loss)

                    # evaluation for each 2 epochs
                    with torch.no_grad():
                        model.eval()
                        auc,logloss = evaluate(task_idx+1,model)
                        print("evaluation auc,logloss on", task_idx+1,"is",auc,logloss)
                        auc_old,logloss_old = evaluate_spec(test_loader_old,model)
                        print("evaluation auc,logloss on", task_idx+1,"is",auc_old,logloss_old)
                        auc_notold,logloss_notold = evaluate_spec(test_loader_notold,model)
                        print("evaluation auc,logloss on", task_idx+1,"is",auc_notold,logloss_notold)
                        if auc_old<global_auc_old or logloss_old>global_logloss_old:
                            if not patience:
                                overall_auc.append(global_auc)
                                overall_logloss.append(global_logloss)
                                overall_auc_old.append(global_auc_old)
                                overall_logloss_old.append(global_logloss_old)
                                overall_auc_notold.append(global_auc_notold)
                                overall_logloss_notold.append(global_logloss_notold)
                                break
                            patience-=1
                        else:
                            patience = global_patience
                            global_auc, global_logloss = auc, logloss
                            global_auc_old, global_logloss_old = auc_old, logloss_old
                            global_auc_notold, global_logloss_notold = auc_notold, logloss_notold

                        if epoch == epoch_num-1:
                            overall_auc.append(global_auc)
                            overall_logloss.append(global_logloss)
                            overall_auc_old.append(global_auc_old)
                            overall_logloss_old.append(global_logloss_old)
                            overall_auc_notold.append(global_auc_notold)
                            overall_logloss_notold.append(global_logloss_notold)
            wandb.log({"train loss":train_loss})
            wandb.log({"auc":global_auc})
            wandb.log({"logloss":global_logloss})

        #logging to wandb
        wandb.log({"train loss":train_loss})
        wandb.log({"auc":sum(overall_auc)/len(overall_auc)})
        wandb.log({"logloss":sum(overall_logloss)/len(overall_logloss)})


__author__ = 'maru'

import numpy as np

def auc_ranked(neutral_target, ranking, pos_label=0):
    class_dist = np.bincount(neutral_target)
    neutral_count = 0.0
    auc_sum = 0.0
    for i in ranking:

        if neutral_target[i] == pos_label:
            neutral_count += 1
        else:
            auc_sum += neutral_count * (1.0 / class_dist[0]) * (1.0 / class_dist[1])
    return auc_sum

def auc_uncertainty(neutral_target, proba_pos, pos_label=0):

    """
    compute AUC values based on uncertainty ranking. This only works for neutral classificaiton of instances
    Note: only works for binary neutral classification
    @param neutral_target: target values of neutral classification(class)
    @param proba_pos: probabilities of positive class
    @param pos_label: index of positive labels
    """
    uncertainty = proba_pos.min(axis=1)
    ranking = np.argsort(uncertainty)[::-1] # descending order of uncertainty

    return auc_ranked(neutral_target, ranking, pos_label=pos_label)


from sys import maxsize #maxsize is the maximum value of a variable of type int
v = 4 #number of vertices

def TSP(graph, s): #s is the starting vertex
    vertex = []
    for i in range(v):
        if i != s:
            vertex.append(i) #vertex is a list of vertices excluding s

    min_path = maxsize #
    while True: #loop until all permutations are checked
        current_cost = 0
        k = s
        for i in range(len(vertex)):
            current_cost += graph[k][vertex[i]]
            k = vertex[i]
        current_cost += graph[k][s]
        min_path = min(min_path, current_cost)

        if not next_perm(vertex):
            break
    return min_path

def next_perm(l): #l is a list of vertices
    n = len(l)
    i = n-2

    while i >= 0 and l[i] > l[i+1]: #find the largest index i such that l[i] < l[i+1]
        i -= 1

    if i == -1: #if no such index exists, the permutation is the last permutation
        return False

    j = i+1
    while j < n and l[j] > l[i]: #find the smallest element greater than l[i]
        j += 1 #j is the index of the smallest element greater than l[i]

    j -= 1 #decrement j to get the index of the smallest element greater than l[i]

    l[i], l[j] = l[j], l[i] #swap l[i] and l[j]
    left = i+1
    right = n-1

    while left < right: #reverse the list
        l[left], l[right] = l[right], l[left]
        left += 1
        right -= 1
    return True

graph = [[0,10,15,20], [10,0,35,25], [15,35,0,30], [20,25,30,0]] #adjacency matrix of the graph
s = 0
res = TSP(graph,s) #s is the starting vertex of the graph and res is the minimum cost of the TSP
print("Minimun cost of TSP:",res) #print the result


"""GIDDY: GeospatIal Distribution DYnamics

Giddy is an open-source python library for the analysis of dynamics of
longitudinal spatial data. Originating from the spatial dynamics module
in PySAL (Python Spatial Analysis Library), it is under active development
for the inclusion of many newly proposed analytics that consider the
role of space in the evolution of distributions over time and has
several new features including inter- and intra-regional decomposition
of mobility association and local measures of exchange mobility in
addition to space-time LISA and spatial markov methods. Give
giddy a try if you are interested in space-time analysis!

"""

DOCLINES = __doc__.split("\n")

from setuptools import setup
from distutils.command.build_py import build_py

setup(name='giddy',  # name of package
      version='1.0.2',
      description=DOCLINES[0],
      long_description="\n".join(DOCLINES[2:]),
      url='https://github.com/pysal/giddy',
      maintainer='Wei Kang',
      maintainer_email='weikang9009@gmail.com',
      test_suite='nose.collector',
      tests_require=['nose'],
      keywords='spatial statistics, spatiotemporal analysis',
      classifiers=[
        'Development Status :: 5 - Production/Stable',
        'Intended Audience :: Science/Research',
        'Intended Audience :: Developers',
        'Intended Audience :: Education',
        'Topic :: Scientific/Engineering',
        'Topic :: Scientific/Engineering :: GIS',
        'License :: OSI Approved :: BSD License',
        'Programming Language :: Python',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3.4',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6'
        ],
      license='3-Clause BSD',
      packages=['giddy'],
      install_requires=['mapclassify', 'esda'],
      zip_safe=False,
      cmdclass={'build.py': build_py})


# Imports
import dash
from dash import dcc, no_update
from dash import html
from dash import dash_table
from dash.dependencies import Input, Output, State
import dash_bootstrap_components as dbc
from dash.dash_table.Format import Format, Scheme
import plotly.graph_objs as go
import threading
import pandas as pd
import warnings

# Relative imports
from database.connector import MongoDBConnector
from analytics import DashTable, BackTest
from utils import get_ticker_name

# Turn off warnings (from optimial parameters)
warnings.filterwarnings("ignore")

# Initialize Dash app
app = dash.Dash(external_stylesheets=[dbc.themes.BOOTSTRAP])
app.title = "Equity Pairs Trading Tool"
app.update_title = None
server = app.server

# Initialize MongoDB connector
db_connector = MongoDBConnector()

####################################
# LAYOUT
####################################
app.layout = html.Div(
    [
        # Header
        html.H1("Equity Pairs Trading Tool"),
        html.H3(
            "1. Select the time range, number of pairs and selection criterion."
        ),
        html.H3(
            "2. Press 'Compute pairs'."
        ),
        html.H3(
            "3. Select a pair to backtest a trading strategy."
        ),
        html.H3(
            "4. Experiment with backtesting parameters."
        ),
        html.H5(
            f"Equities included: individual components (and indices) of the S&P 500, NASDAQ 100, and Russell 2000. Total number of equities included: {len(db_connector.tickers_in_db)}"
        ),
        html.H6("Notes: (i) Data is only available from 1 Jan 2020 to 31 Jul 2023."),
        html.H6(
            "(ii) In order to reduce computation time, only pairs with above 0.95 correlation are considered."
        ),
        html.Br(),
        # Date Picker for start date
        html.Label("Start Date:"),
        dcc.DatePickerSingle(id="start-date-picker", date="2022-01-01"),
        # Date Picker for end date
        html.Label("End Date:"),
        dcc.DatePickerSingle(id="end-date-picker", date="2022-12-31"),
        # Number of pairs
        html.Label("Number of Pairs:"),
        dcc.Input(
            id="num_of_pairs",
            type="number",
            value=100,  # Default value
            min=1,
            step=1,
        ),
        html.Br(),
        # Dropdown for criteria selection
        html.Label("Criterion for Pairs Selection:"),
        dcc.Dropdown(
            id="pairs-criteria-dropdown",
            options=[
                {"label": "Correlation", "value": "Correlation"},
                {
                    "label": "Mean Reversion Half Life",
                    "value": "Mean Reversion Half Life",
                },
                {"label": "Cointegration P-Value", "value": "Cointegration P-Value"},
            ],
            value="Cointegration P-Value",  # Default value
            clearable=False,  # Prevents the user from clearing the selection
            style={"width": "50%"},
        ),
        # Break
        html.Br(),
        # Button to compute correlations
        html.Button("Compute pairs", id="submit-button", n_clicks=0),
        # Error message
        html.Div(id="input-error-message"),
        # Interval component to update the display
        dcc.Store(id="process-state"),
        dcc.Interval(
            id="interval-component",
            interval=1e12,  # some massive number, essentially inf
            n_intervals=0,
            max_intervals=1,  # Only run once
        ),
        # Display area for results (tables, graphs, etc.)
        html.Div(
            [
                html.Div(
                    [  # This div will contain the table and be on the left
                        html.Div(id="display-results"),
                    ],
                    style={
                        "width": "50%",
                        "display": "inline-block",
                        "vertical-align": "top",
                    },
                ),
                html.Div(
                    [
                        dcc.Graph(id="price-plot"),
                    ],
                    style={
                        "width": "50%",
                        "display": "inline-block",
                        "vertical-align": "top",
                    },
                ),
            ],
            style={"display": "flex", "flex-direction": "row"},
        ),
        html.H2("Backtesting Strategy"),
        html.H5(
            "A simple mean reversion strategy is employed, based on z-scores. The z-score is calculated as: (spread - rolling mean of spread) / rolling standard deviation of spread. If the z-score is above the threshold, then a long position in the cheaper stock and a short position in the more expensive stock are taken. If the z-score is below the negative of the threshold, then a short position in the cheaper stock and a long position in the more expensive stock are taken."
        ),
        html.H5(
            "You can adjust the threshold and rolling window below, and press Find Optimal Parameters to see the parameters that maximise the Sharpe Ratio."
        ),
        html.H5(
            "A z-score threshold of 1.0 and a rolling window of 30 days are used by default."
        ),
        html.H6(
            "Notes (i): The risk-free rate has been assumed to be the Effective Fed Funds Rate (EFFR)."
        ),
        html.H6(
            "(ii) 'Find Optimal Parameters' is a rudimentary optimisation method (Bayesian Optimisation), and may not be the true optimal parameters. Optimal parameters may be different every time you compute optimal parameters."
        ),
        # Options for backtesting
        html.Div(
            [
                html.Label("z-score Threshold:"),
                dcc.Input(id="z-threshold", type="number", value=1.0, step=0.01),
                html.Label("Rolling Window:"),
                dcc.Input(id="rolling_window", type="number", value=30, step=0.01),
            ]
        ),
        # Optimise button
        html.Button("Find Optimal Parameters", id="optimize-button", n_clicks=0),
        html.Div(id="optimal-params-display"),
        html.Div(
            [html.Div(id="backtest-table"), dcc.Graph(id="strategy-performance-plot")],
            style={
                "width": "50%",
                "display": "inline-block",
                "vertical-align": "top",
            },
        ),
        # Backtesting graphs
        html.Div(
            [
                dcc.Graph(id="backtest-graph"),
            ],
            style={
                "width": "50%",
                "display": "inline-block",
                "vertical-align": "top",
            },
        ),
    ]
)

####################################
# MAIN TABLE
####################################
# Global variable to store the computed data
done_event = threading.Event()
computed_data = None

def threaded_task(
    start_date,
    end_date,
    num_of_pairs,
    pairs_criteria=0.95,
    selection_criteria="Correlation",
):
    """
    Function to run in the background to compute the data for the table.
    :param start_date: The start date.
    :param end_date: The end date.
    :param num_of_pairs: The number of pairs to return in the table
    :param pairs_criteria: The minimum correlation considered
    :param selection_criteria: The criteria for selecting the pairs.
    """
    global computed_data
    analytics = DashTable(start_date, end_date, num_of_pairs, pairs_criteria)
    # Limit down to only where correlation is above criteria
    analytics.compute_top_correlations()
    # Compute data for the table
    computed_data = analytics.stock_selection(selection_criteria)
    done_event.set()

# CALLBACK TO UPDATE THE TABLE
@app.callback(
    [
        Output("display-results", "children"),
        Output("input-error-message", "children"),
        Output("process-state", "data"),
        Output("interval-component", "interval"),
        Output("interval-component", "max_intervals"),
    ],  # Add this output
    [Input("submit-button", "n_clicks"), Input("interval-component", "n_intervals")],
    [Input("pairs-criteria-dropdown", "value")],
    [
        State("start-date-picker", "date"),
        State("end-date-picker", "date"),
        State("num_of_pairs", "value"),
        State("process-state", "data"),
    ],
)
def update_display(
    n_clicks,
    n_intervals,
    selected_criteria,
    start_date,
    end_date,
    num_of_pairs,
    process_state,
):
    ctx = dash.callback_context
    if not ctx.triggered:
        triggered_id = None
    else:
        triggered_id = ctx.triggered[0]["prop_id"].split(".")[0]

    # START THE BACKGROUND TASK AND SHOW A SPINNER
    if triggered_id == "submit-button":
        # If the button is clicked, reset the done_event and clear the previously computed data
        done_event.clear()
        global computed_data
        computed_data = None

        # Start the background task and immediately return a spinner
        thread = threading.Thread(
            target=threaded_task,
            args=(start_date, end_date, num_of_pairs, 0.95, selected_criteria),
        )
        thread.start()
        spinner = dbc.Spinner(size="lg", color="primary")
        return (
            spinner,
            "",
            "processing",
            1000,
            1e12,
        )  # Start checking every second and let it run indefinitely

    # WHEN BACKGROUND TASK IS DONE, DISPLAY THE RESULTS
    elif triggered_id == "interval-component" and done_event.is_set():
        data, cols = computed_data
        number_format = Format(precision=3, scheme=Scheme.fixed)
        formatted_cols = [
            {**col, "type": "numeric", "format": number_format} for col in cols
        ]
        table = dash_table.DataTable(
            data=data,
            columns=formatted_cols,
            page_size=10,
            id="table",
            row_selectable="single",
            sort_action="native",
        )
        return (
            table,
            "",
            "done",
            1e12,
            n_intervals,
        )  # Stop the interval component by setting max_intervals to its current value

    # IF NO UPDATE, DO NOTHING
    else:
        return no_update, no_update, no_update, no_update, no_update  # No updates

####################################
# PAIRS GRAPH
####################################
@app.callback(
    Output("price-plot", "figure"),
    [
        Input("table", "selected_rows"),
        Input("start-date-picker", "date"),
        Input("end-date-picker", "date"),
    ],
    [State("table", "data")],
)
def update_graph(selected_rows, start_date, end_date, table_data):
    if not selected_rows:  # If no rows are selected, return an empty figure
        return {}

    # Get the selected pair from the table data
    selected_pair = table_data[selected_rows[0]]["Pair"]
    stock1, stock2 = selected_pair.split(" - ")

    # Get name of the stocks
    stock1_name = get_ticker_name(stock1)
    stock2_name = get_ticker_name(stock2)

    # Fetch the price data for the selected stocks within the date range
    stock_1_price = db_connector.fetch_single_stock_on_time_range(
        stock1, start_date, end_date
    )
    stock_2_price = db_connector.fetch_single_stock_on_time_range(
        stock2, start_date, end_date
    )

    # Create the plot
    fig = {
        "data": [
            {
                "x": stock_1_price.index,
                "y": stock_1_price["close"],
                "type": "line",
                "name": stock1,
            },
            {
                "x": stock_2_price.index,
                "y": stock_2_price["close"],
                "type": "line",
                "name": stock2,
            },
        ],
        "layout": {
            "title": f"Price Data for {selected_pair}
{stock1_name} and {stock2_name}",
            "labels": {"x": "Date", "y": "Price (USD)"},
        },
    }
    return fig

####################################
# BACKTEST TABLE AND SPREAD GRAPH
####################################
# BACKTESTING
@app.callback(
    [Output("backtest-graph", "figure"), Output("backtest-table", "children")],
    [
        Input("table", "selected_rows"),
        Input("start-date-picker", "date"),
        Input("end-date-picker", "date"),
        Input("z-threshold", "value"),
        Input("rolling_window", "value"),
    ],
    [State("table", "data")],
)
def update_backtest_graph(
    selected_rows, start_date, end_date, z_thresh, rolling_window, table_data
):
    if not selected_rows:
        return {}

    # Get the selected pair from the table data
    selected_pair = table_data[selected_rows[0]]["Pair"]
    stock1, stock2 = selected_pair.split(" - ")

    # Run backtest
    backtest = BackTest(stock1, stock2, start_date, end_date)
    backtest.backtest(window=rolling_window, z_thresh=z_thresh)
    signal = backtest.signal
    (
        sharpe_ratio,
        sortino_ratio,
        max_drawdown,
        annualized_return,
        annualized_volatility,
    ) = backtest.calculate_metrics()

    # Create the table
    data_dict = {
        "Metric": [
            "Sharpe Ratio",
            "Sortino Ratio",
            "Max Drawdown",
            "Annualized Return",
            "Annualized Volatility",
        ],
        "Value": [
            sharpe_ratio,
            sortino_ratio,
            max_drawdown,
            annualized_return,
            annualized_volatility,
        ],
    }
    df = pd.DataFrame(data_dict)
    data = df.to_dict("records")

    metrics = {
        "Sharpe Ratio": sharpe_ratio,
        "Sortino Ratio": sortino_ratio,
        "Max Drawdown": max_drawdown,
        "Annualized Return": annualized_return,
        "Annualized Volatility": annualized_volatility,
    }

    formatted_metrics = {}
    for key, value in metrics.items():
        if key in ["Max Drawdown", "Annualized Return", "Annualized Volatility"]:
            formatted_metrics[key] = "{:.2%}".format(value)
        else:
            formatted_metrics[key] = "{:.2f}".format(value)

    data_dict = {
        "Metric": list(metrics.keys()),
        "Value": list(formatted_metrics.values()),
    }
    df = pd.DataFrame(data_dict)
    data = df.to_dict("records")

    columns = [{"name": "Metric", "id": "Metric"}, {"name": "Value", "id": "Value"}]

    table = dash_table.DataTable(data=data, columns=columns)

    # Create the plot
    fig = go.Figure()

    # Add the main line plot for spread
    fig.add_trace(
        go.Scatter(
            x=backtest.signal.index, y=backtest.spread, mode="lines", name="Spread"
        )
    )

    # Add background shading based on signal
    shapes = []
    dates = signal.index.tolist()
    signal = signal.tolist()
    for i, sig in enumerate(signal):
        if i == len(signal) - 1:
            continue

        color = None
        if sig == 1:
            color = "green"
        elif sig == -1:
            color = "red"

        shapes.append(
            {
                "type": "rect",
                "xref": "x",
                "yref": "paper",
                "x0": dates[i],
                "y0": 0,
                "x1": dates[i + 1],
                "y1": 1,
                "fillcolor": color,
                "opacity": 0.2,
                "line_width": 0,
            }
        )

    fig.update_layout(
        shapes=shapes,
        title=f"Spread for {selected_pair}
Green and red regions represent long and short positions in the spread",
        plot_bgcolor="white",
        paper_bgcolor="white",
        title_font_color="black",
        xaxis=dict(
            title="Date",
            titlefont=dict(color="black"),
            tickfont=dict(color="black"),
            gridcolor="lightgrey",
        ),
        yaxis=dict(
            title="Spread (USD)",
            titlefont=dict(color="black"),
            tickfont=dict(color="black"),
            gridcolor="lightgrey",
        ),
    )
    return fig, table

####################################
# PRINT OUT OPTIMAL PARAMETERS
####################################
@app.callback(
    Output("optimal-params-display", "children"),
    [Input("optimize-button", "n_clicks")],
    [
        State("start-date-picker", "date"),
        State("end-date-picker", "date"),
        State("table", "selected_rows"),
        State("table", "data"),
    ],
)
def find_optimal_params(n_clicks, start_date, end_date, selected_rows, table_data):
    if not selected_rows or n_clicks == 0:
        return dash.no_update

    # Get the selected pair from the table data
    selected_pair = table_data[selected_rows[0]]["Pair"]
    stock1, stock2 = selected_pair.split(" - ")

    backtest_optim = BackTest(stock1, stock2, start_date, end_date)
    optimal_params = backtest_optim.find_optimal_params()

    # Format and display the parameters
    return html.P(
        f"Parameters to optimise the Sharpe Ratio: z-score: {optimal_params['z_thresh']}, rolling window: {optimal_params['window']}"
    )

####################################
# CUMULATATIVE RETURNS PLOT
####################################
@app.callback(
    Output("strategy-performance-plot", "figure"),
    [
        Input("table", "selected_rows"),
        Input("z-threshold", "value"),
        Input("rolling_window", "value"),
    ],
    [
        State("start-date-picker", "date"),
        State("end-date-picker", "date"),
        State("table", "data"),
    ],
)
def update_strategy_performance(
    selected_rows, z_threshold, rolling_window, start_date, end_date, table_data
):
    if not selected_rows:
        return dash.no_update

    # Get the selected pair from the table data
    selected_pair = table_data[selected_rows[0]]["Pair"]
    stock1, stock2 = selected_pair.split(" - ")

    backtest = BackTest(stock1, stock2, start_date, end_date)
    backtest.backtest(window=rolling_window, z_thresh=z_threshold)

    fig = {
        "data": [
            {
                "x": backtest.cumulative_returns.index,
                "y": backtest.cumulative_returns,
                "type": "line",
            },
        ],
        "layout": {
            "title": f"Cumulatative Returns",
            "labels": {"x": "Date", "y": "Returns"},
        },
    }

    return fig

if __name__ == "__main__":
    app.run_server(debug=False)


"""
Given a string containing only three types of characters: '(', ')' and '*', write a function to check whether this string is valid. We define the validity of a string by these rules:

Any left parenthesis '(' must have a corresponding right parenthesis ')'.
Any right parenthesis ')' must have a corresponding left parenthesis '('.
Left parenthesis '(' must go before the corresponding right parenthesis ')'.
'*' could be treated as a single right parenthesis ')' or a single left parenthesis '(' or an empty string.
An empty string is also valid.
Example 1:
Input: "()"
Output: True
Example 2:
Input: "(*)"
Output: True
Example 3:
Input: "(*))"
Output: True
Note:
The string size will be in the range [1, 100].
"""

class Solution:
    def checkValidString(self, instr: str) -> bool:
        oStack, sStack, cStack = [],[],[]
        for i,ch in enumerate(instr):
            if ch == '*': sStack.append(i)
            elif ch == '(': oStack.append(i)
            else:
                if oStack and oStack[-1] < i:
                    oStack.pop()
                else:
                    cStack.append(i)

        cStack = cStack[::-1]
        sStack = sStack[::-1]
        while cStack:
            if not sStack: break
            if sStack[-1] < cStack[-1]:
                sStack.pop()
                cStack.pop()
            else:
                break

        sStack = sStack[::-1]
        while oStack:
            if not sStack: break
            if sStack[-1] > oStack[-1]:
                sStack.pop()
                oStack.pop()
            else:
                break
        return not (oStack or cStack)

s = Solution()
print(s.checkValidString("((*)*"))

Leetcode/Leetcode - Premium/Google/Valid Parenthesis String.py


from flask import Flask, request
from flask.views import MethodView
from flask_cors import CORS

import assistance as asst

app = Flask(__name__)
CORS(app)

class PhChaserAPI(MethodView):

    def get(self):
        code = request.args.get("CODE")
        phone = request.args.get("PHONE")
        return asst.Tracker.take_by_code(code, phone)

app.add_url_rule('/PhChaser', view_func=PhChaserAPI.as_view('PhChaser'))

class PyCountriesAPI(MethodView):

    def get(self, iso):
        if not iso:
            return asst.PyCountries.take()
        # Country Code
        return asst.PyCountries.take_code_by(iso)

PyCountriesVIEW = PyCountriesAPI.as_view('PyCountries')
app.add_url_rule('/PyCountries', view_func=PyCountriesVIEW, defaults={'iso': None})
app.add_url_rule('/PyCountries/<string:iso>/Code', view_func=PyCountriesVIEW,
                 methods=['GET', 'PUT', 'DELETE'])

class PyCountriesCodeAPI(MethodView):

    def get(self):
        return asst.PyCountries.take_code()

PyCountriesCodeVIEW = PyCountriesCodeAPI.as_view('PyCountriesCode')
app.add_url_rule('/PyCountriesCode', view_func=PyCountriesCodeVIEW)

if __name__ == '__main__':
    app.run(debug=True)


import sys
import os
import astropy.io.fits as pyfits
import numpy

from podi_definitions import *

import scipy.interpolate

def subtract_pupilghost(filename_hdu, output_file, profile_file, operation="subtract", verbose=True):

    # Make the routine flexible in handling both a filename or HDUList as input parameter
    if (type(filename_hdu) is str):
        if (os.path.isfile(fitsfile)):
            hdulist = pyfits.open(fitsfile)
        else:
            return
    else:
        hdulist = filename_hdu

    filter = hdulist[0].header['FILTER']

    if (os.path.isfile(profile_file) and not operation=="ignore"):
        profile_data = numpy.loadtxt(profile_file)

        # Loop over all extensions
        # For now only use the first one, hard enough
        for ota_id in range(1, len(hdulist)):

            ota = hdulist[ota_id]
            extname = ota.header["EXTNAME"]

            if (filter in pupilghost_centers):
                if (extname in pupilghost_centers[filter]):
                    center_x, center_y = pupilghost_centers[filter][extname]

                    if (verbose): stdout_write("\rCorrecting pupil-ghost in extension %s ..." % (extname))
                    # Create the pixel array and create all distances
                    #        buffer = numpy.zeros(shape=(ota.data.shape[1], ota.data.shape[0]))
                    x, y = numpy.indices(ota.data.shape)
                    dx = x - center_x
                    dy = y - center_y
                    d = numpy.sqrt(dx*dx + dy*dy)

                    # Next load the profile, and fit a spline that we can use to evaluate it
                    profile = scipy.interpolate.interp1d(profile_data[:,0]*4, profile_data[:,1], bounds_error=False, fill_value=0)
                    fit = profile(d.ravel())

                    if (operation == "replace"):
                        hdulist[ota_id].data = fit.reshape(ota.data.shape).transpose()
                    else:
                        hdulist[ota_id].data -= fit.reshape(ota.data.shape).transpose()

                    ota.header.add_history("Pupilghost: profile in %s" % profile_file)
                    ota.header.add_history("Pupilghost: center @ %d,%d" % (center_x, center_y))

    if (output_file is not None):
        # Now all the work is done, all final data is stored in img2, write results to new file
        if (verbose): stdout_write(" writing output ...")
        clobberfile(outputfile)
        hdulist.writeto(outputfile)
        if (verbose): stdout_write(" done!\n\n")
    else:
        if (verbose): stdout_write(" done, handing on data!\n\n")
        return hdulist

if __name__ == "__main__":

    # Read in the input parameters
    fitsfile = sys.argv[1]
    profile_file = sys.argv[2]
    outputfile = sys.argv[3]

    operation = cmdline_arg_set_or_default("-op", "subtract")

    subtract_pupilghost(fitsfile, outputfile, profile_file, operation=operation, verbose=True)


#
import glob
import os
import argparse

from PIL import Image
import matplotlib.pyplot as plt

import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

from sklearn.metrics import roc_curve, roc_auc_score, auc
from sklearn.model_selection import  GroupKFold

from models import *
from bilinear_layers import *

import utils

class oaDataset(Dataset):
    def __init__(self, root_path, transform):
        # store filenames
        self.roi_root = root_path
        self.flist = glob.glob(os.path.join(self.roi_root, "**"))
        self.transform = transform

    def __len__(self):
        # return size of dataset
        return len(self.flist)

    def __getitem__(self, idx):
        fname = self.flist[idx]
        image, min_jsw, jsw_des, fjsw, kl, jsize = np.load(fname, allow_pickle=True)  # PIL image

        if kl > 1:
            grade =1
        else:
            grade = 0

        image = Image.fromarray(image.astype('uint8'), 'L')
        image = self.transform(image)

        jsw = np.array(jsw_des)
        # noise = np.random.normal(0, 1, 221)
        # jsw = jsw + noise
        jsw = torch.from_numpy(jsw).float()

        return image, jsw, grade

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--oai_assesment', default='/../DATA/OAI/kXR_SQ_BU00_SAS/kxr_sq_bu00.sas7bdat')
    parser.add_argument('--model', default=combined)
    args = parser.parse_args()

    SEED = 42
    MAX_EPOCH = 100

    np.random.seed(SEED)
    torch.manual_seed(SEED);
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if use_cuda else "cpu")
    cudnn.benchmark = True

    cv = GroupKFold(n_splits=5)

    train_transforms = transforms.Compose([
        transforms.Resize(size=(56,56)),
        transforms.RandomCrop(size= (48,48)),
        transforms.ToTensor(),
    ])

    eval_transforms = transforms.Compose([
        transforms.Resize(size=(48,48)),
        transforms.ToTensor(),
    ])

    plt.figure(figsize=(10, 10))

    train_root = '../DATA/CROPS/OAI_00m_tm_fjsw_standardized/'
    test_root = '../DATA/CROPS/MOST_00m_tm_fjsw_standardized/'

    train_dataloader = DataLoader(oaDataset(train_root, train_transforms),
                                  batch_size=64,
                                  shuffle=True,
                                  num_workers=8,
                                  drop_last=True)
    test_dataloader = DataLoader(oaDataset(test_root, eval_transforms),
                                 batch_size=64,
                                 num_workers=8,
                                 drop_last=True)

    net = args.model().to(device)

    # Loss and optimizer
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, nesterov=True)

    best_acc = 0.0
    best_epoch = None

    fpr = []
    tpr = []

    for epoch in range(MAX_EPOCH):

        optimizer,lr = utils.adjust_learning_rate(optimizer,epoch,0.00001, 0.01, 8)

        # Training
        train_loss = utils.train_epoch(epoch, net, optimizer, train_dataloader, criterion)
        # Validating
        val_loss, preds, truth = utils.validate_epoch(net, test_dataloader, criterion)

        auc_val = roc_auc_score(truth, preds)
        print(epoch + 1, train_loss, val_loss, auc_val)

        if auc_val > best_acc:
            best_acc = auc_val
            best_epoch = epoch
            fpr, tpr, thresholds = roc_curve(truth, preds)

    plt.plot(fpr, tpr, lw=2, alpha=0.8,
             label='ROC (AUC = %0.2f)' % (best_acc))
    plt.show()

    print("Best AUC:", best_acc)


from tensorflow.keras.layers import Input, Dense, BatchNormalization,Conv1D,Activation,MaxPooling1D,Concatenate, Flatten
# from tensorflow.keras import backend as K
# from tensorflow.keras.models import Model
from keras.engine import Layer
from keras import initializers
from keras import backend as K

from keras.layers import Dense, Input, Flatten
from keras.layers import Conv1D, MaxPooling1D, GRU, Bidirectional, TimeDistributed
from keras.models import Model

def recall_m(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall = true_positives / (possible_positives + K.epsilon())
    return recall

def precision_m(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision = true_positives / (predicted_positives + K.epsilon())
    return precision

def f1_m(y_true, y_pred):
    precision = precision_m(y_true, y_pred)
    recall = recall_m(y_true, y_pred)
    return 2 * ((precision * recall) / (precision + recall + K.epsilon()))

class AttLayer(Layer):
    def __init__(self, attention_dim):
        self.init = initializers.get('normal')
        self.supports_masking = True
        self.attention_dim = attention_dim
        super(AttLayer, self).__init__()

    def build(self, input_shape):
        assert len(input_shape) == 3
        self.W = K.variable(self.init((input_shape[-1], self.attention_dim)), name='W')
        self.b = K.variable(self.init((self.attention_dim, )), name='b')
        self.u = K.variable(self.init((self.attention_dim, 1)), name='u')
        self.trainable_weights = [self.W, self.b, self.u]
        super(AttLayer, self).build(input_shape)

    def compute_mask(self, inputs, mask=None):
        return mask

    def call(self, x, mask=None):
        # size of x :[batch_size, sel_len, attention_dim]
        # size of u :[batch_size, attention_dim]
        # uit = tanh(xW+b)
        uit = K.tanh(K.bias_add(K.dot(x, self.W), self.b))
        ait = K.dot(uit, self.u)
        ait = K.squeeze(ait, -1)

        ait = K.exp(ait)

        if mask is not None:
            # Cast the mask to floatX to avoid float64 upcasting in theano
            ait *= K.cast(mask, K.floatx())
        ait /= K.cast(K.sum(ait, axis=1, keepdims=True) + K.epsilon(), K.floatx())
        ait = K.expand_dims(ait)
        weighted_input = x * ait
        output = K.sum(weighted_input, axis=1)

        return output

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[-1])

class Ye_CNN():
    def __init__(self):
        self.model = None

    def cnn_model(self, tok_len, embed_size):
        filter_sizes = [1, 2, 3, 4, 5, 6, 7]
        num_filters = 128
        sent_input = Input(shape=(tok_len, embed_size),dtype='float')
        conv_0 = Conv1D(num_filters, kernel_size=(filter_sizes[0]))(sent_input)
        act_0 = Activation('relu')(conv_0)
        bn_0 = BatchNormalization()(act_0)

        conv_1 = Conv1D(num_filters, kernel_size=(filter_sizes[1]))(sent_input)
        act_1 = Activation('relu')(conv_1)
        bn_1 = BatchNormalization()(act_1)

        conv_2 = Conv1D(num_filters, kernel_size=(filter_sizes[2]))(sent_input)
        act_2 = Activation('relu')(conv_2)
        bn_2 = BatchNormalization()(act_2)

        conv_3 = Conv1D(num_filters, kernel_size=(filter_sizes[3]))(sent_input)
        act_3 = Activation('relu')(conv_3)
        bn_3 = BatchNormalization()(act_3)

        conv_4 = Conv1D(num_filters, kernel_size=(filter_sizes[4]))(sent_input)
        act_4 = Activation('relu')(conv_4)
        bn_4 = BatchNormalization()(act_4)

        conv_5 = Conv1D(num_filters, kernel_size=(filter_sizes[5]))(sent_input)
        act_5 = Activation('relu')(conv_5)
        bn_5 = BatchNormalization()(act_5)

        conv_6 = Conv1D(num_filters, kernel_size=(filter_sizes[6]))(sent_input)
        act_6 = Activation('relu')(conv_6)
        bn_6 = BatchNormalization()(act_6)

        maxpool_0 = MaxPooling1D(pool_size=(tok_len - filter_sizes[0]))(bn_0)
        maxpool_1 = MaxPooling1D(pool_size=(tok_len - filter_sizes[1]))(bn_1)
        maxpool_2 = MaxPooling1D(pool_size=(tok_len - filter_sizes[2]))(bn_2)
        maxpool_3 = MaxPooling1D(pool_size=(tok_len - filter_sizes[3]))(bn_3)
        maxpool_4 = MaxPooling1D(pool_size=(tok_len - filter_sizes[4]))(bn_4)
        maxpool_5 = MaxPooling1D(pool_size=(tok_len - filter_sizes[5]))(bn_5)
        maxpool_6 = MaxPooling1D(pool_size=(tok_len - filter_sizes[6]))(bn_6)

        sent_concat = Concatenate()([maxpool_0, maxpool_1, maxpool_2, maxpool_3, maxpool_4,
                                     maxpool_5, maxpool_6])
        sent_flat = Flatten()(sent_concat)
        sent_den = Dense(units=100, activation='relu')(sent_flat)
        pred = Dense(12, activation='softmax')(sent_den)

        self.model = Model(sent_input, pred)
        self.model.summary()
        self.model.compile(loss='categorical_crossentropy', metrics=['acc', precision_m, recall_m, f1_m],
                           optimizer='adam')

    def build_model(self, max_tok, embed_size):
        self.cnn_model(max_tok, embed_size)

    def train(self, x_train, y_train, x_val, y_val, batch_size, checkpoint=None):
        self.model.fit(x_train, y_train, epochs=30, batch_size=batch_size, validation_data=[x_val, y_val],
                       verbose=1, callbacks=[checkpoint])

        loss, accuracy, precision, recall, f1_score = self.model.evaluate(x_val, y_val, verbose=1)

        return loss, accuracy, precision, recall, f1_score

class Ye_RNN_ATT():
    def __init__(self):
        self.model = None

    def rnn_model(self, tok_len, embed_size):
        sent_input = Input(shape=(tok_len, embed_size), dtype='float')
        sent_lstm = Bidirectional(GRU(100, return_sequences=True))(sent_input)
        att = AttLayer(100)(sent_lstm)
        pred = Dense(12, activation='softmax')(att)
        self.model = Model(sent_input, pred)
        self.model.summary()
        self.model.compile(loss='categorical_crossentropy', metrics=['acc', precision_m, recall_m, f1_m],
                           optimizer='adam')

    def build_model(self, max_tok, embed_size):
        self.rnn_model(max_tok, embed_size)

    def train(self, x_train, y_train, x_val, y_val, batch_size, checkpoint=None):
        self.model.fit(x_train, y_train, epochs=30, batch_size=batch_size, validation_data=[x_val, y_val],
                       verbose=1, callbacks=[checkpoint])

        loss, accuracy, precision, recall, f1_score = self.model.evaluate(x_val, y_val, verbose=1)

        return loss, accuracy, precision, recall, f1_score

class Ye_HAN():
    def __init__(self):
        self.model = None

    def han_model(self, seq_len, tok_len, embed_size):
        sentence_input = Input(shape=(seq_len, embed_size), dtype='float')
        l_lstm = Bidirectional(GRU(100, return_sequences=True))(sentence_input)
        l_att = AttLayer(100)(l_lstm)
        sentEncoder = Model(sentence_input, l_att)

        review_input = Input(shape=(seq_len, tok_len, embed_size), dtype='float')
        review_encoder = TimeDistributed(sentEncoder)(review_input)
        l_lstm_sent = Bidirectional(GRU(100, return_sequences=True))(review_encoder)
        l_att_sent = AttLayer(100)(l_lstm_sent)
        preds = Dense(11, activation='softmax')(l_att_sent)
        self.model = Model(review_input, preds)

        self.model.compile(loss='categorical_crossentropy',
                      optimizer='adam',
                      metrics=['acc', precision_m, recall_m, f1_m])

    def build_model(self,seq_len, max_tok, embed_size):
        self.han_model(seq_len, max_tok, embed_size)

    def train(self, x_train, y_train, x_val, y_val, batch_size, checkpoint=None):
        self.model.fit(x_train, y_train, epochs=30, batch_size=batch_size, validation_data=[x_val, y_val],
                       verbose=1, callbacks=[checkpoint])

        loss, accuracy, precision, recall, f1_score = self.model.evaluate(x_val, y_val, verbose=1)

        return loss, accuracy, precision, recall, f1_score