#!/usr/bin/env python3
# lib/gdb.py

# internal imports
import html
import logging
import re
from datetime import datetime

# python <3.9 workaround
try:
    from zoneinfo import ZoneInfo
except ImportError:
    from backports.zoneinfo import ZoneInfo

# external imports
import mechanicalsoup
from netaddr import IPNetwork

# local imports

# module logger
log = logging.getLogger(__name__)

class GRAYDB:
    """Class for interacting with the gray database"""

    def __init__(self, url, credentials):
        """Init of graydb class

        url: string         - base url to database

        credentials: dict {
            login: str      - user name
            password: str   - password
        }
        """
        self.browser = mechanicalsoup.StatefulBrowser(
            raise_on_404=True,
            user_agent='inkotools-api/0.4',
        )
        self.baseurl = url
        self.credentials = credentials
        self._login()

    def __del__(self):
        self.browser.close()

    class CredentialsError(Exception):
        """Custom exception on wrong creds"""

        def __init__(self, msg="Wrong login or password!"):
            self.message = msg
            super().__init__(self.message)

    class NotFoundError(Exception):
        """Custom exception for not found errors"""

        def __init__(self, msg="Client not found"):
            self.message = msg
            super().__init__(self.message)

    def _login(self):
        """Login to graydb"""
        b = self.browser
        b.open(self.baseurl)
        # check that there is the auth form
        if len(b.page.select('form[name=auth]')) > 0:
            b.select_form('form[name=auth]')
            b['username'] = self.credentials['login'].encode('cp1251')
            b['password'] = self.credentials['password'].encode('cp1251')
            b.submit_selected()
            # check if auth form again - wrong login
            if len(b.page.select('form[name=auth]')) > 0:
                raise self.CredentialsError()
            else:
                log.debug('Logged in successfully')
        else:
            log.debug('Already logged in')

    def _check_auth(func):
        """Check auth decorator"""

        def wrapper(self, *args, **kwargs):
            self._login()
            return func(self, *args, **kwargs)
        return wrapper

    def get_client_ip_list(self, contract_id: str):
        """Get list of client ips from billing"""
        inet = self.get_billing_accounts(contract_id)['internet']
        if len(inet) == 0:
            raise self.NotFoundError('Internet account not found')
        return inet['ip_list']

    def get_billing_accounts(self, contract_id: str):
        """Get list of client services from billing"""
        raw = self.browser.post(f'{self.baseurl}/bil.php',
                                data={"nome_dogo": contract_id})
        account_types = [
            "internet",
            "telephony",
            "ld_telephony",
            "television",
        ]
        # init res dict with empty accounts
        res = dict.fromkeys(account_types, {})
        # iterate through all accounts (first row is table header)
        for idx, row in enumerate(raw.soup.select('table tr')[1:]):
            item = {}
            item['account_id'] = int(row.contents[0].string)
            item['services'] = list(map(gdb_decode, row.contents[1].strings))
            # remove `,` and `;` symbols and split without empty strings
            ip_tel = row.contents[3].string.translate(
                {ord(i): None for i in ';,'}).split()
            # internet
            if idx == 0:
                item['tariff'] = gdb_decode(row.contents[2].string)
                item['ip_list'] = ip_tel
            # telephony
            elif idx < 3:
                item['number_list'] = ip_tel
            item['balance'] = float(row.contents[5].string)
            item['credit'] = float(row.contents[7].string)
            status = gdb_decode(row.contents[8].string)
            item['enabled'] = True if status == 'Разблокирован' else False
            res[account_types[idx]] = item

        return res

    def get_billing_speed(self, ip: str):
        """Get billing speed for ip address"""
        ip = str(ip)
        raw = self.browser.get(f'http://62.182.48.36/speed/index.php?ip1={ip}')
        raw = raw.soup.get_text().strip()
        rgx = r'esteblished= ([.\d]+ Mbit/sec)'
        r = re.findall(rgx, raw)
        if len(r) == 0:
            res = "error"
            log.error(raw)
        else:
            res = r[0]
        return res

    def get_billing_summary(self, contract_id: str):
        """Get billing accounts with billing speed included"""
        res = self.get_billing_accounts(contract_id)
        if res['internet'] and len(res['internet']['ip_list']) > 0:
            # speed is the same for all ips, so, select first
            res['internet']['speed'] = self.get_billing_speed(
                res['internet']['ip_list'][0])
        return res

    def get_contract_by_ip(self, client_ip: str):
        """Find contract with client ip"""
        client_ip = str(client_ip)
        raw = self.browser.post(f'{self.baseurl}/poisk_test.php',
                                data={"ip": client_ip, "go99": 1})
        # graydb has fuzzy search, so iterate through several contracts
        # and check if contract's ip is the searched ip
        for row in raw.soup.select('tbody tr'):
            contract_id = row.find('td').string.strip()
            # workaround for billing random fails
            for i in range(3):
                try:
                    client_ips = self.get_client_ip_list(contract_id)
                    break
                except self.NotFoundError:
                    log.warning(f'Failed! Attempt: {i+1}/3')
            if client_ip in client_ips:
                return contract_id
        raise self.NotFoundError(f"Client with ip '{client_ip}' not found")

    def get_free_ip(self, vid: int, return_list: bool = False):
        """Find free ip in billing"""
        base = f'192.168.{vid}.'
        res = []
        # search in ip range splitting them by 10
        for i in range(1, 13):
            expr = base+str(i)
            raw = self.browser.post(f'{self.baseurl}/poisk_test.php',
                                    data={"ip": expr, "go99": 1})
            cnt = len(raw.soup.select('tbody tr'))
            # cnt < 11 --> there are free ips
            # cnt > 11 and i in [1,2] --> 1xx and 2xx ips workaround
            if cnt != 11:
                for j in ['']+list(range(0, 10)):
                    ip = expr+str(j)
                    host = int(ip.split('.')[3])
                    if host == 1:
                        continue
                    if host > 128:
                        break
                    try:
                        self.get_contract_by_ip(ip)
                    except self.NotFoundError:
                        log.debug(f'found: {ip}')
                        if return_list:
                            res.append(ip)
                        else:
                            return ip
        if len(res) > 0:
            return res
        raise self.NotFoundError(f'No free ip addresses found in vlan {vid}')

    def get_vlan_ip_count(self, vid: int):
        """Get ip count in vlan from billing"""
        expr = f'192.168.{vid}.'
        raw = self.browser.post(f'{self.baseurl}/poisk_test.php',
                                data={"ip": expr, "go99": 1})
        cnt = len(raw.soup.select('tbody tr'))
        return cnt

    def get_free_pip(self, subnet: str, return_list: bool = False):
        """Find free public ip in billing"""
        res = []
        subnet = IPNetwork(subnet)
        for ip in list(map(str, subnet.iter_hosts()))[1::]:
            try:
                self.get_contract_by_ip(ip)
            except self.NotFoundError:
                log.debug(f'found: {ip}')
                if return_list:
                    res.append(ip)
                else:
                    return ip
        if len(res) > 0:
            return res
        raise self.NotFoundError(f'No free ip addresses found for {subnet}')

    def get_internal_client_id(self, contract_id: str):
        """Get internal client id in gray database"""
        raw = self.browser.post(f'{self.baseurl}/poisk_test.php',
                                data={"dogovor": contract_id, "startt": 1})
        f = raw.soup.find('input', {'name': 'id_aabon'})
        if f is None:
            raise self.NotFoundError(f"Contract '{contract_id}' not found")
        res = int(f.get('value'))
        return res

    @_check_auth
    def get_client_data(self, contract_id: str):
        """Get client info from gray database"""
        client_id = self.get_internal_client_id(contract_id)
        raw = self.browser.get(f'{self.baseurl}/index.php',
                               params={"id_aabon": client_id})
        raw = raw.soup
        res = {'contract_id': contract_id, 'client_id': client_id}
        # matching dict between returning keys and form input names
        m_dict = {
            'name': 'fio',
            'company': 'organizatsiya',
            'house': 'dom',
            'room': 'kvartira',
            'office': 'ofis_tseh',
            'sw_ip': 'loyalnost',
            'port': 'port',
            'cable_length': 'dlina_cab',
        }
        # select form with client data (form with input `fio`)
        d = list(raw.find('input', {'name': 'fio'}).parents)[6]
        # iterate through the form inputs
        for key, val in m_dict.items():
            res[key] = d.find(attrs={'name': val}).get('value').strip()
        # get city and street from first (selected) option in select
        m_dict = {
            'city': 'gorod',
            'street': 'ulitsa',
            'entrance': 'podezd',
            'floor': 'etaj',
        }
        for key, val in m_dict.items():
            res[key] = d.find(
                'select', {'name': val}).option.get('value').strip()
        # generate contact list
        res['contact_list'] = []
        for i in range(1, 4):
            c = d.find(attrs={'name': f'cont{i}'}).get('value').strip()
            # skip empty strings and legacy values
            if c in ['', '0']:
                continue
            # ignore duplicates
            if c not in res['contact_list']:
                res['contact_list'].append(c)
        # comment string from textarea
        res['comment'] = d.find(attrs={'name': 'primechanie'}).string
        res['comment'] = res['comment'].strip(
        ) if res['comment'] is not None else ''
        # search for terminated mark
        res['terminated'] = bool(
            raw.find('font', {'color': 'red', 'size': '2px'}))
        # clear legacy values from old db
        for i in res:
            if res[i] == '0':
                res[i] = ''
        # search for tickets
        tickets = []
        keys = ['ticket_id', 'issue', 'date', 'master', 'creator', 'comments']
        for t in raw.find_all(attrs={'name': 'id_start_zay'}):
            # get columns as reverted previous siblings of current
            cols = t.parent.find_previous_siblings()[4::-1]
            # generate values from td tags without last (comments)
            values = [x.string.split(':', 1)[1].strip() for x in cols[:-1]]
            # split comments (last column)
            comments = parse_comments(cols[-1])
            # join ticket_id and other values
            values = [int(t['value']), *values, comments]
            ticket = dict(zip(keys, values))
            # convert date from string to date
            ticket['date'] = datetime.strptime(
                ticket['date'], '%d-%m-%y').replace(
                tzinfo=ZoneInfo('Europe/Moscow'))
            # raw comments for add_comment
            ticket['raw_comments'] = t.parent.find(
                attrs={"name": "id_start_stadya"})["value"]
            tickets.append(ticket)
        res['tickets'] = tickets
        return res

    @_check_auth
    def get_tickets(self):
        """Get list of user tickets"""
        tickets = []
        keys = ['ticket_id', 'contract_id', 'name', 'issue', 'address',
                'contacts', 'date', 'creator']
        raw = self.browser.get(f'{self.baseurl}/zayavki.php')
        raw = raw.soup
        for row in raw.tbody.find_all('tr'):
            values = list(map(lambda x: re.sub(
                r' +', ' ', gdb_decode(x.text).strip()), row.find_all('td')))
            # normal user tickets
            if len(values) == 11:
                values = values[0:8]
            # boss tickets
            elif len(values) == 13:
                values = values[1:9]
            # error
            else:
                self.log.error(f'Wrong ticket structure: {values}')
                continue
            ticket = dict(zip(keys, values))
            # strip '№ ' from ticket id and convert to int
            ticket['ticket_id'] = int(ticket['ticket_id'][2:])
            # convert contacts to list and remove duplicates
            ticket['contacts'] = list(set(ticket['contacts'].split()))
            # convert str to date
            ticket['date'] = datetime.strptime(
                ticket['date'], '%d-%m-%y').replace(
                tzinfo=ZoneInfo('Europe/Moscow'))
            # add comments
            ticket['comments'] = parse_comments(row)
            tickets.append(ticket)
        return tickets

    @_check_auth
    def add_comment(self, contract_id: str, ticket_id: int, comment: str):
        """Add comment to ticket"""
        ticket_found = False
        client = self.get_client_data(contract_id)
        # ticket validation and getting old comments
        # in graydb we need to add all the old comments with the new one
        for i in client["tickets"]:
            if i["ticket_id"] == ticket_id:
                ticket_found = True
                old_comments = i["raw_comments"]
                break
        if not ticket_found:
            raise self.NotFoundError('Ticket not found')
        comment = html.escape(comment)
        data = {"tekst_zay": comment.encode("cp1251"),
                "id_start_zay": ticket_id,
                "id_start_stadya": old_comments.encode("cp1251"),
                }
        user = self.credentials['login']
        log.debug(f'[{user}] contract {contract_id}, ticket {ticket_id}, '
                  f'comment: {comment}')
        res = self.browser.post(f'{self.baseurl}/index.php', data=data)
        if res.ok:
            log.info(f'[{user}] commented [{contract_id}]')
            return 'Comment added successfully'
        else:
            log.error(f'[{user}] failed to comment [{contract_id}]')
            return {'error': 'Failed to add comment'}

    @_check_auth
    def search_ticket(self, contract_id: str, keywords: list = [],
                      search_in_comments: bool = False,
                      master: str = ''):
        """Search ticket by keywords"""
        res = []
        if not isinstance(keywords, list):
            keywords = [keywords]
        client = self.get_client_data(contract_id)
        kw = r'(?i)' + r'|'.join(keywords)
        for ticket in client['tickets']:
            if ((master == '' or ticket['master'] == master)
                and (re.search(kw, ticket['issue'])
                     or (search_in_comments
                         and re.search(kw, ticket['raw_comments'])))):
                res.append(ticket)
        return res

    @_check_auth
    def change_client_data(self, contract_id: str, data: dict):
        """Change client data in gray database

        Possible data: sw_ip, port, cable_length

        Returns: bool"""
        m_dict = {
            'sw_ip': 'loyalnost',
            'port': 'port',
            'cable_length': 'dlina_cab',
        }
        client_id = self.get_internal_client_id(contract_id)
        # get old user data from gray database
        raw = self.browser.get(f'{self.baseurl}/index.php',
                               params={"id_aabon": client_id})
        raw = raw.soup
        form = list(raw.find('input', {'name': 'fio'}).parents)[6]
        fields = form.find_all(['input', 'select', 'textarea'])
        old_data = {}
        for item in fields:
            if (not 'name' in item.attrs
                    or item.attrs['name'] in ['zayavka', 'otvetstv']):
                # skip submit button and ticket inputs
                continue
            key = item.attrs['name']
            if item.name == 'input':
                val = item.attrs['value'].encode("cp1251")
            elif item.name == 'textarea':
                val = item.text.encode("cp1251")
            elif item.name == 'select':
                val = item.option.get('value').encode("cp1251")
            old_data[key] = val
        # update data with provided values
        new_data = old_data
        for k, v in data.items():
            if k in m_dict:
                new_data[m_dict[k]] = v
        # send data to gray database
        raw = self.browser.post(f'{self.baseurl}/index.php', data=new_data)
        res = list(raw.soup.stripped_strings)[-1] == 'Завершено успешно!'
        if res:
            log.debug(f'[{contract_id}] data changed: {data}')
        else:
            log.error(f'[{contract_id}] failed to change data: {data}')
        return res

    @_check_auth
    def terminate_contract(self, contract_id: str):
        """Terminate contract in gray database"""
        # check for services in billing
        for k, v in self.get_billing_accounts(contract_id).items():
            if 'services' in v.keys() and len(v['services']) > 0:
                log.warning(
                    f'Found services in {k} billing account, aborting.')
                return False
        # send termination request
        self.browser.post(f'{self.baseurl}/index.php',
                          data={"rastorg": contract_id})
        return True

########################################################################
# common functions

def parse_comments(tag):
    """Search for comments within soup tag"""
    comments = []
    for i, e in enumerate(tag.find_all(href='#')):
        comments.append({
            "time": datetime.strptime(
                e['title'], '%d.%m.%Y %H:%M:%S').replace(
                tzinfo=ZoneInfo('Europe/Moscow')),
            "author": gdb_decode(e.string),
            "comment": list(map(
                gdb_decode, e.parent.strings))[(i+1)*2].strip()[2:-1],
        })
    return comments

def gdb_decode(s):
    try:
        res = s.encode('ISO-8859-1').decode('cp1251')
    except Exception:
        res = s
    return res


import re
import os
import time
import argparse
from os import walk

def removeFromName(dir_path, reg):
    dirpath, dirnames, filenames = next(walk(dir_path))
    for directory in dirnames:
        removeFromName(directory, reg)

    for filename in filenames:
        newName = re.sub(reg, '', filename).strip()
        if(newName == ""): continue
        print(newName)
        os.rename(dirpath + "/" + filename, dirpath + "/" + newName)

parser = argparse.ArgumentParser(description='Convert to HEVC video encoder.')
parser.add_argument(
        '-r', '--regex', required=True, type=str, help='regex')
parser.add_argument(
        '-d', '--directory', required=True, type=str, help='directory path')

args = parser.parse_args()
regex = args.regex
dir_path = args.directory

removeFromName(dir_path, regex)


# A test suite for pdb; not very comprehensive at the moment.

import doctest
import os
import pdb
import sys
import types
import unittest
import subprocess
import textwrap

from test import support
# This little helper class is essential for testing pdb under doctest.
from test.test_doctest import _FakeInput

class PdbTestInput(object):
    """Context manager that makes testing Pdb in doctests easier."""

    def __init__(self, input):
        self.input = input

    def __enter__(self):
        self.real_stdin = sys.stdin
        sys.stdin = _FakeInput(self.input)
        self.orig_trace = sys.gettrace() if hasattr(sys, 'gettrace') else None

    def __exit__(self, *exc):
        sys.stdin = self.real_stdin
        if self.orig_trace:
            sys.settrace(self.orig_trace)

def test_pdb_displayhook():
    """This tests the custom displayhook for pdb.

    >>> def test_function(foo, bar):
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     pass

    >>> with PdbTestInput([
    ...     'foo',
    ...     'bar',
    ...     'for i in range(5): print(i)',
    ...     'continue',
    ... ]):
    ...     test_function(1, None)
    > <doctest test.test_pdb.test_pdb_displayhook[0]>(3)test_function()
    -> pass
    (Pdb) foo
    1
    (Pdb) bar
    (Pdb) for i in range(5): print(i)
    0
    1
    2
    3
    4
    (Pdb) continue
    """

def test_pdb_basic_commands():
    """Test the basic commands of pdb.

    >>> def test_function_2(foo, bar='default'):
    ...     print(foo)
    ...     for i in range(5):
    ...         print(i)
    ...     print(bar)
    ...     for i in range(10):
    ...         never_executed
    ...     print('after for')
    ...     print('...')
    ...     return foo.upper()

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     ret = test_function_2('baz')
    ...     print(ret)

    >>> with PdbTestInput([  # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    ...     'step',       # entering the function call
    ...     'args',       # display function args
    ...     'list',       # list function source
    ...     'bt',         # display backtrace
    ...     'up',         # step up to test_function()
    ...     'down',       # step down to test_function_2() again
    ...     'next',       # stepping to print(foo)
    ...     'next',       # stepping to the for loop
    ...     'step',       # stepping into the for loop
    ...     'until',      # continuing until out of the for loop
    ...     'next',       # executing the print(bar)
    ...     'jump 8',     # jump over second for loop
    ...     'return',     # return out of function
    ...     'retval',     # display return value
    ...     'continue',
    ... ]):
    ...    test_function()
    > <doctest test.test_pdb.test_pdb_basic_commands[1]>(3)test_function()
    -> ret = test_function_2('baz')
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(1)test_function_2()
    -> def test_function_2(foo, bar='default'):
    (Pdb) args
    foo = 'baz'
    bar = 'default'
    (Pdb) list
      1  ->     def test_function_2(foo, bar='default'):
      2             print(foo)
      3             for i in range(5):
      4                 print(i)
      5             print(bar)
      6             for i in range(10):
      7                 never_executed
      8             print('after for')
      9             print('...')
     10             return foo.upper()
    [EOF]
    (Pdb) bt
    ...
      <doctest test.test_pdb.test_pdb_basic_commands[2]>(18)<module>()
    -> test_function()
      <doctest test.test_pdb.test_pdb_basic_commands[1]>(3)test_function()
    -> ret = test_function_2('baz')
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(1)test_function_2()
    -> def test_function_2(foo, bar='default'):
    (Pdb) up
    > <doctest test.test_pdb.test_pdb_basic_commands[1]>(3)test_function()
    -> ret = test_function_2('baz')
    (Pdb) down
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(1)test_function_2()
    -> def test_function_2(foo, bar='default'):
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(2)test_function_2()
    -> print(foo)
    (Pdb) next
    baz
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(3)test_function_2()
    -> for i in range(5):
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(4)test_function_2()
    -> print(i)
    (Pdb) until
    0
    1
    2
    3
    4
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(5)test_function_2()
    -> print(bar)
    (Pdb) next
    default
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(6)test_function_2()
    -> for i in range(10):
    (Pdb) jump 8
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(8)test_function_2()
    -> print('after for')
    (Pdb) return
    after for
    ...
    --Return--
    > <doctest test.test_pdb.test_pdb_basic_commands[0]>(10)test_function_2()->'BAZ'
    -> return foo.upper()
    (Pdb) retval
    'BAZ'
    (Pdb) continue
    BAZ
    """

def test_pdb_breakpoint_commands():
    """Test basic commands related to breakpoints.

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     print(1)
    ...     print(2)
    ...     print(3)
    ...     print(4)

    First, need to clear bdb state that might be left over from previous tests.
    Otherwise, the new breakpoints might get assigned different numbers.

    >>> from bdb import Breakpoint
    >>> Breakpoint.next = 1
    >>> Breakpoint.bplist = {}
    >>> Breakpoint.bpbynumber = [None]

    Now test the breakpoint commands.  NORMALIZE_WHITESPACE is needed because
    the breakpoint list outputs a tab for the "stop only" and "ignore next"
    lines, which we don't want to put in here.

    >>> with PdbTestInput([  # doctest: +NORMALIZE_WHITESPACE
    ...     'break 3',
    ...     'disable 1',
    ...     'ignore 1 10',
    ...     'condition 1 1 < 2',
    ...     'break 4',
    ...     'break 4',
    ...     'break',
    ...     'clear 3',
    ...     'break',
    ...     'condition 1',
    ...     'enable 1',
    ...     'clear 1',
    ...     'commands 2',
    ...     'p "42"',
    ...     'print("42", 7*6)',     # Issue 18764 (not about breakpoints)
    ...     'end',
    ...     'continue',  # will stop at breakpoint 2 (line 4)
    ...     'clear',     # clear all!
    ...     'y',
    ...     'tbreak 5',
    ...     'continue',  # will stop at temporary breakpoint
    ...     'break',     # make sure breakpoint is gone
    ...     'continue',
    ... ]):
    ...    test_function()
    > <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>(3)test_function()
    -> print(1)
    (Pdb) break 3
    Breakpoint 1 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
    (Pdb) disable 1
    Disabled breakpoint 1 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
    (Pdb) ignore 1 10
    Will ignore next 10 crossings of breakpoint 1.
    (Pdb) condition 1 1 < 2
    New condition set for breakpoint 1.
    (Pdb) break 4
    Breakpoint 2 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) break 4
    Breakpoint 3 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) break
    Num Type         Disp Enb   Where
    1   breakpoint   keep no    at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
            stop only if 1 < 2
            ignore next 10 hits
    2   breakpoint   keep yes   at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    3   breakpoint   keep yes   at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) clear 3
    Deleted breakpoint 3 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) break
    Num Type         Disp Enb   Where
    1   breakpoint   keep no    at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
            stop only if 1 < 2
            ignore next 10 hits
    2   breakpoint   keep yes   at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) condition 1
    Breakpoint 1 is now unconditional.
    (Pdb) enable 1
    Enabled breakpoint 1 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
    (Pdb) clear 1
    Deleted breakpoint 1 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:3
    (Pdb) commands 2
    (com) p "42"
    (com) print("42", 7*6)
    (com) end
    (Pdb) continue
    1
    '42'
    42 42
    > <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>(4)test_function()
    -> print(2)
    (Pdb) clear
    Clear all breaks? y
    Deleted breakpoint 2 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:4
    (Pdb) tbreak 5
    Breakpoint 4 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:5
    (Pdb) continue
    2
    Deleted breakpoint 4 at <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>:5
    > <doctest test.test_pdb.test_pdb_breakpoint_commands[0]>(5)test_function()
    -> print(3)
    (Pdb) break
    (Pdb) continue
    3
    4
    """

def do_nothing():
    pass

def do_something():
    print(42)

def test_list_commands():
    """Test the list and source commands of pdb.

    >>> def test_function_2(foo):
    ...     import test.test_pdb
    ...     test.test_pdb.do_nothing()
    ...     'some...'
    ...     'more...'
    ...     'code...'
    ...     'to...'
    ...     'make...'
    ...     'a...'
    ...     'long...'
    ...     'listing...'
    ...     'useful...'
    ...     '...'
    ...     '...'
    ...     return foo

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     ret = test_function_2('baz')

    >>> with PdbTestInput([  # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    ...     'list',      # list first function
    ...     'step',      # step into second function
    ...     'list',      # list second function
    ...     'list',      # continue listing to EOF
    ...     'list 1,3',  # list specific lines
    ...     'list x',    # invalid argument
    ...     'next',      # step to import
    ...     'next',      # step over import
    ...     'step',      # step into do_nothing
    ...     'longlist',  # list all lines
    ...     'source do_something',  # list all lines of function
    ...     'source fooxxx',        # something that doesn't exit
    ...     'continue',
    ... ]):
    ...    test_function()
    > <doctest test.test_pdb.test_list_commands[1]>(3)test_function()
    -> ret = test_function_2('baz')
    (Pdb) list
      1         def test_function():
      2             import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
      3  ->         ret = test_function_2('baz')
    [EOF]
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_list_commands[0]>(1)test_function_2()
    -> def test_function_2(foo):
    (Pdb) list
      1  ->     def test_function_2(foo):
      2             import test.test_pdb
      3             test.test_pdb.do_nothing()
      4             'some...'
      5             'more...'
      6             'code...'
      7             'to...'
      8             'make...'
      9             'a...'
     10             'long...'
     11             'listing...'
    (Pdb) list
     12             'useful...'
     13             '...'
     14             '...'
     15             return foo
    [EOF]
    (Pdb) list 1,3
      1  ->     def test_function_2(foo):
      2             import test.test_pdb
      3             test.test_pdb.do_nothing()
    (Pdb) list x
    *** ...
    (Pdb) next
    > <doctest test.test_pdb.test_list_commands[0]>(2)test_function_2()
    -> import test.test_pdb
    (Pdb) next
    > <doctest test.test_pdb.test_list_commands[0]>(3)test_function_2()
    -> test.test_pdb.do_nothing()
    (Pdb) step
    --Call--
    > ...test_pdb.py(...)do_nothing()
    -> def do_nothing():
    (Pdb) longlist
    ...  ->     def do_nothing():
    ...             pass
    (Pdb) source do_something
    ...         def do_something():
    ...             print(42)
    (Pdb) source fooxxx
    *** ...
    (Pdb) continue
    """

def test_post_mortem():
    """Test post mortem traceback debugging.

    >>> def test_function_2():
    ...     try:
    ...         1/0
    ...     finally:
    ...         print('Exception!')

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     test_function_2()
    ...     print('Not reached.')

    >>> with PdbTestInput([  # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    ...     'next',      # step over exception-raising call
    ...     'bt',        # get a backtrace
    ...     'list',      # list code of test_function()
    ...     'down',      # step into test_function_2()
    ...     'list',      # list code of test_function_2()
    ...     'continue',
    ... ]):
    ...    try:
    ...        test_function()
    ...    except ZeroDivisionError:
    ...        print('Correctly reraised.')
    > <doctest test.test_pdb.test_post_mortem[1]>(3)test_function()
    -> test_function_2()
    (Pdb) next
    Exception!
    ZeroDivisionError: division by zero
    > <doctest test.test_pdb.test_post_mortem[1]>(3)test_function()
    -> test_function_2()
    (Pdb) bt
    ...
      <doctest test.test_pdb.test_post_mortem[2]>(10)<module>()
    -> test_function()
    > <doctest test.test_pdb.test_post_mortem[1]>(3)test_function()
    -> test_function_2()
      <doctest test.test_pdb.test_post_mortem[0]>(3)test_function_2()
    -> 1/0
    (Pdb) list
      1         def test_function():
      2             import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
      3  ->         test_function_2()
      4             print('Not reached.')
    [EOF]
    (Pdb) down
    > <doctest test.test_pdb.test_post_mortem[0]>(3)test_function_2()
    -> 1/0
    (Pdb) list
      1         def test_function_2():
      2             try:
      3  >>             1/0
      4             finally:
      5  ->             print('Exception!')
    [EOF]
    (Pdb) continue
    Correctly reraised.
    """

def test_pdb_skip_modules():
    """This illustrates the simple case of module skipping.

    >>> def skip_module():
    ...     import string
    ...     import pdb; pdb.Pdb(skip=['stri*'], nosigint=True, readrc=False).set_trace()
    ...     string.capwords('FOO')

    >>> with PdbTestInput([
    ...     'step',
    ...     'continue',
    ... ]):
    ...     skip_module()
    > <doctest test.test_pdb.test_pdb_skip_modules[0]>(4)skip_module()
    -> string.capwords('FOO')
    (Pdb) step
    --Return--
    > <doctest test.test_pdb.test_pdb_skip_modules[0]>(4)skip_module()->None
    -> string.capwords('FOO')
    (Pdb) continue
    """

# Module for testing skipping of module that makes a callback
mod = types.ModuleType('module_to_skip')
exec('def foo_pony(callback): x = 1; callback(); return None', mod.__dict__)

def test_pdb_skip_modules_with_callback():
    """This illustrates skipping of modules that call into other code.

    >>> def skip_module():
    ...     def callback():
    ...         return None
    ...     import pdb; pdb.Pdb(skip=['module_to_skip*'], nosigint=True, readrc=False).set_trace()
    ...     mod.foo_pony(callback)

    >>> with PdbTestInput([
    ...     'step',
    ...     'step',
    ...     'step',
    ...     'step',
    ...     'step',
    ...     'continue',
    ... ]):
    ...     skip_module()
    ...     pass  # provides something to "step" to
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(5)skip_module()
    -> mod.foo_pony(callback)
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(2)callback()
    -> def callback():
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(3)callback()
    -> return None
    (Pdb) step
    --Return--
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(3)callback()->None
    -> return None
    (Pdb) step
    --Return--
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(5)skip_module()->None
    -> mod.foo_pony(callback)
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[1]>(10)<module>()
    -> pass  # provides something to "step" to
    (Pdb) continue
    """

def test_pdb_continue_in_bottomframe():
    """Test that "continue" and "next" work properly in bottom frame (issue #5294).

    >>> def test_function():
    ...     import pdb, sys; inst = pdb.Pdb(nosigint=True, readrc=False)
    ...     inst.set_trace()
    ...     inst.botframe = sys._getframe()  # hackery to get the right botframe
    ...     print(1)
    ...     print(2)
    ...     print(3)
    ...     print(4)

    >>> with PdbTestInput([  # doctest: +ELLIPSIS
    ...     'next',
    ...     'break 7',
    ...     'continue',
    ...     'next',
    ...     'continue',
    ...     'continue',
    ... ]):
    ...    test_function()
    > <doctest test.test_pdb.test_pdb_continue_in_bottomframe[0]>(4)test_function()
    -> inst.botframe = sys._getframe()  # hackery to get the right botframe
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_continue_in_bottomframe[0]>(5)test_function()
    -> print(1)
    (Pdb) break 7
    Breakpoint ... at <doctest test.test_pdb.test_pdb_continue_in_bottomframe[0]>:7
    (Pdb) continue
    1
    2
    > <doctest test.test_pdb.test_pdb_continue_in_bottomframe[0]>(7)test_function()
    -> print(3)
    (Pdb) next
    3
    > <doctest test.test_pdb.test_pdb_continue_in_bottomframe[0]>(8)test_function()
    -> print(4)
    (Pdb) continue
    4
    """

def pdb_invoke(method, arg):
    """Run pdb.method(arg)."""
    getattr(pdb.Pdb(nosigint=True, readrc=False), method)(arg)

def test_pdb_run_with_incorrect_argument():
    """Testing run and runeval with incorrect first argument.

    >>> pti = PdbTestInput(['continue',])
    >>> with pti:
    ...     pdb_invoke('run', lambda x: x)
    Traceback (most recent call last):
    TypeError: exec() arg 1 must be a string, bytes or code object

    >>> with pti:
    ...     pdb_invoke('runeval', lambda x: x)
    Traceback (most recent call last):
    TypeError: eval() arg 1 must be a string, bytes or code object
    """

def test_pdb_run_with_code_object():
    """Testing run and runeval with code object as a first argument.

    >>> with PdbTestInput(['step','x', 'continue']):  # doctest: +ELLIPSIS
    ...     pdb_invoke('run', compile('x=1', '<string>', 'exec'))
    > <string>(1)<module>()...
    (Pdb) step
    --Return--
    > <string>(1)<module>()->None
    (Pdb) x
    1
    (Pdb) continue

    >>> with PdbTestInput(['x', 'continue']):
    ...     x=0
    ...     pdb_invoke('runeval', compile('x+1', '<string>', 'eval'))
    > <string>(1)<module>()->None
    (Pdb) x
    1
    (Pdb) continue
    """

def test_next_until_return_at_return_event():
    """Test that pdb stops after a next/until/return issued at a return debug event.

    >>> def test_function_2():
    ...     x = 1
    ...     x = 2

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     test_function_2()
    ...     test_function_2()
    ...     test_function_2()
    ...     end = 1

    >>> from bdb import Breakpoint
    >>> Breakpoint.next = 1
    >>> with PdbTestInput(['break test_function_2',
    ...                    'continue',
    ...                    'return',
    ...                    'next',
    ...                    'continue',
    ...                    'return',
    ...                    'until',
    ...                    'continue',
    ...                    'return',
    ...                    'return',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(3)test_function()
    -> test_function_2()
    (Pdb) break test_function_2
    Breakpoint 1 at <doctest test.test_pdb.test_next_until_return_at_return_event[0]>:1
    (Pdb) continue
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
    -> x = 1
    (Pdb) return
    --Return--
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
    -> x = 2
    (Pdb) next
    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(4)test_function()
    -> test_function_2()
    (Pdb) continue
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
    -> x = 1
    (Pdb) return
    --Return--
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
    -> x = 2
    (Pdb) until
    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(5)test_function()
    -> test_function_2()
    (Pdb) continue
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
    -> x = 1
    (Pdb) return
    --Return--
    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
    -> x = 2
    (Pdb) return
    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(6)test_function()
    -> end = 1
    (Pdb) continue
    """

def test_pdb_next_command_for_generator():
    """Testing skip unwindng stack on yield for generators for "next" command

    >>> def test_gen():
    ...     yield 0
    ...     return 1
    ...     yield 2

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     it = test_gen()
    ...     try:
    ...         if next(it) != 0:
    ...             raise AssertionError
    ...         next(it)
    ...     except StopIteration as ex:
    ...         if ex.value != 1:
    ...             raise AssertionError
    ...     print("finished")

    >>> with PdbTestInput(['step',
    ...                    'step',
    ...                    'step',
    ...                    'next',
    ...                    'next',
    ...                    'step',
    ...                    'step',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(3)test_function()
    -> it = test_gen()
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(4)test_function()
    -> try:
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(5)test_function()
    -> if next(it) != 0:
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(1)test_gen()
    -> def test_gen():
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(2)test_gen()
    -> yield 0
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(3)test_gen()
    -> return 1
    (Pdb) step
    --Return--
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(3)test_gen()->1
    -> return 1
    (Pdb) step
    StopIteration: 1
    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(7)test_function()
    -> next(it)
    (Pdb) continue
    finished
    """

if False:
    def test_pdb_next_command_for_coroutine():
        """Testing skip unwindng stack on yield for coroutines for "next" command

        >>> import asyncio

        >>> async def test_coro():
        ...     await asyncio.sleep(0)
        ...     await asyncio.sleep(0)
        ...     await asyncio.sleep(0)

        >>> async def test_main():
        ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
        ...     await test_coro()

        >>> def test_function():
        ...     loop = asyncio.new_event_loop()
        ...     loop.run_until_complete(test_main())
        ...     loop.close()
        ...     print("finished")

        >>> with PdbTestInput(['step',
        ...                    'step',
        ...                    'next',
        ...                    'next',
        ...                    'next',
        ...                    'step',
        ...                    'continue']):
        ...     test_function()
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[2]>(3)test_main()
        -> await test_coro()
        (Pdb) step
        --Call--
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[1]>(1)test_coro()
        -> async def test_coro():
        (Pdb) step
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[1]>(2)test_coro()
        -> await asyncio.sleep(0)
        (Pdb) next
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[1]>(3)test_coro()
        -> await asyncio.sleep(0)
        (Pdb) next
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[1]>(4)test_coro()
        -> await asyncio.sleep(0)
        (Pdb) next
        Internal StopIteration
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[2]>(3)test_main()
        -> await test_coro()
        (Pdb) step
        --Return--
        > <doctest test.test_pdb.test_pdb_next_command_for_coroutine[2]>(3)test_main()->None
        -> await test_coro()
        (Pdb) continue
        finished
        """

    def test_pdb_next_command_for_asyncgen():
        """Testing skip unwindng stack on yield for coroutines for "next" command

        >>> import asyncio

        >>> async def agen():
        ...     yield 1
        ...     await asyncio.sleep(0)
        ...     yield 2

        >>> async def test_coro():
        ...     async for x in agen():
        ...         print(x)

        >>> async def test_main():
        ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
        ...     await test_coro()

        >>> def test_function():
        ...     loop = asyncio.new_event_loop()
        ...     loop.run_until_complete(test_main())
        ...     loop.close()
        ...     print("finished")

        >>> with PdbTestInput(['step',
        ...                    'step',
        ...                    'next',
        ...                    'next',
        ...                    'step',
        ...                    'next',
        ...                    'continue']):
        ...     test_function()
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[3]>(3)test_main()
        -> await test_coro()
        (Pdb) step
        --Call--
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[2]>(1)test_coro()
        -> async def test_coro():
        (Pdb) step
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[2]>(2)test_coro()
        -> async for x in agen():
        (Pdb) next
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[2]>(3)test_coro()
        -> print(x)
        (Pdb) next
        1
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[2]>(2)test_coro()
        -> async for x in agen():
        (Pdb) step
        --Call--
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[1]>(2)agen()
        -> yield 1
        (Pdb) next
        > <doctest test.test_pdb.test_pdb_next_command_for_asyncgen[1]>(3)agen()
        -> await asyncio.sleep(0)
        (Pdb) continue
        2
        finished
        """

    def test_pdb_return_command_for_coroutine():
        """Testing no unwindng stack on yield for coroutines for "return" command

    >>> import asyncio

    >>> async def test_coro():
    ...     await asyncio.sleep(0)
    ...     await asyncio.sleep(0)
    ...     await asyncio.sleep(0)

    >>> async def test_main():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     await test_coro()

    >>> def test_function():
    ...     loop = asyncio.new_event_loop()
    ...     loop.run_until_complete(test_main())
    ...     loop.close()
    ...     print("finished")

    >>> with PdbTestInput(['step',
    ...                    'step',
    ...                    'next',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_return_command_for_coroutine[2]>(3)test_main()
    -> await test_coro()
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_return_command_for_coroutine[1]>(1)test_coro()
    -> async def test_coro():
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_return_command_for_coroutine[1]>(2)test_coro()
    -> await asyncio.sleep(0)
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_return_command_for_coroutine[1]>(3)test_coro()
    -> await asyncio.sleep(0)
    (Pdb) continue
    finished
    """

    def test_pdb_until_command_for_coroutine():
        """Testing no unwindng stack for coroutines
           for "until" command if target breakpoing is not reached

        >>> import asyncio

        >>> async def test_coro():
        ...     print(0)
        ...     await asyncio.sleep(0)
        ...     print(1)
        ...     await asyncio.sleep(0)
        ...     print(2)
        ...     await asyncio.sleep(0)
        ...     print(3)

        >>> async def test_main():
        ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
        ...     await test_coro()

        >>> def test_function():
        ...     loop = asyncio.new_event_loop()
        ...     loop.run_until_complete(test_main())
        ...     loop.close()
        ...     print("finished")

        >>> with PdbTestInput(['step',
        ...                    'until 8',
        ...                    'continue']):
        ...     test_function()
        > <doctest test.test_pdb.test_pdb_until_command_for_coroutine[2]>(3)test_main()
        -> await test_coro()
        (Pdb) step
        --Call--
        > <doctest test.test_pdb.test_pdb_until_command_for_coroutine[1]>(1)test_coro()
        -> async def test_coro():
        (Pdb) until 8
        0
        1
        2
        > <doctest test.test_pdb.test_pdb_until_command_for_coroutine[1]>(8)test_coro()
        -> print(3)
        (Pdb) continue
        3
        finished
        """

def test_pdb_return_command_for_generator():
    """Testing no unwindng stack on yield for generators
       for "return" command

    >>> def test_gen():
    ...     yield 0
    ...     return 1
    ...     yield 2

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     it = test_gen()
    ...     try:
    ...         if next(it) != 0:
    ...             raise AssertionError
    ...         next(it)
    ...     except StopIteration as ex:
    ...         if ex.value != 1:
    ...             raise AssertionError
    ...     print("finished")

    >>> with PdbTestInput(['step',
    ...                    'step',
    ...                    'step',
    ...                    'return',
    ...                    'step',
    ...                    'step',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(3)test_function()
    -> it = test_gen()
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(4)test_function()
    -> try:
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(5)test_function()
    -> if next(it) != 0:
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[0]>(1)test_gen()
    -> def test_gen():
    (Pdb) return
    StopIteration: 1
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(7)test_function()
    -> next(it)
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(8)test_function()
    -> except StopIteration as ex:
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(9)test_function()
    -> if ex.value != 1:
    (Pdb) continue
    finished
    """

def test_pdb_until_command_for_generator():
    """Testing no unwindng stack on yield for generators
       for "until" command if target breakpoing is not reached

    >>> def test_gen():
    ...     yield 0
    ...     yield 1
    ...     yield 2

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     for i in test_gen():
    ...         print(i)
    ...     print("finished")

    >>> with PdbTestInput(['step',
    ...                    'until 4',
    ...                    'step',
    ...                    'step',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_until_command_for_generator[1]>(3)test_function()
    -> for i in test_gen():
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(1)test_gen()
    -> def test_gen():
    (Pdb) until 4
    0
    1
    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(4)test_gen()
    -> yield 2
    (Pdb) step
    --Return--
    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(4)test_gen()->2
    -> yield 2
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_until_command_for_generator[1]>(4)test_function()
    -> print(i)
    (Pdb) continue
    2
    finished
    """

def test_pdb_next_command_in_generator_for_loop():
    """The next command on returning from a generator controlled by a for loop.

    >>> def test_gen():
    ...     yield 0
    ...     return 1

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     for i in test_gen():
    ...         print('value', i)
    ...     x = 123

    >>> with PdbTestInput(['break test_gen',
    ...                    'continue',
    ...                    'next',
    ...                    'next',
    ...                    'next',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(3)test_function()
    -> for i in test_gen():
    (Pdb) break test_gen
    Breakpoint 6 at <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>:1
    (Pdb) continue
    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>(2)test_gen()
    -> yield 0
    (Pdb) next
    value 0
    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>(3)test_gen()
    -> return 1
    (Pdb) next
    Internal StopIteration: 1
    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(3)test_function()
    -> for i in test_gen():
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(5)test_function()
    -> x = 123
    (Pdb) continue
    """

def test_pdb_next_command_subiterator():
    """The next command in a generator with a subiterator.

    >>> def test_subgenerator():
    ...     yield 0
    ...     return 1

    >>> def test_gen():
    ...     x = yield from test_subgenerator()
    ...     return x

    >>> def test_function():
    ...     import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
    ...     for i in test_gen():
    ...         print('value', i)
    ...     x = 123

    >>> with PdbTestInput(['step',
    ...                    'step',
    ...                    'next',
    ...                    'next',
    ...                    'next',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(3)test_function()
    -> for i in test_gen():
    (Pdb) step
    --Call--
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(1)test_gen()
    -> def test_gen():
    (Pdb) step
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(2)test_gen()
    -> x = yield from test_subgenerator()
    (Pdb) next
    value 0
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(3)test_gen()
    -> return x
    (Pdb) next
    Internal StopIteration: 1
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(3)test_function()
    -> for i in test_gen():
    (Pdb) next
    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(5)test_function()
    -> x = 123
    (Pdb) continue
    """

def test_pdb_issue_20766():
    """Test for reference leaks when the SIGINT handler is set.

    >>> def test_function():
    ...     i = 1
    ...     while i <= 2:
    ...         sess = pdb.Pdb()
    ...         sess.set_trace(sys._getframe())
    ...         print('pdb %d: %s' % (i, sess._previous_sigint_handler))
    ...         i += 1

    >>> with PdbTestInput(['continue',
    ...                    'continue']):
    ...     test_function()
    > <doctest test.test_pdb.test_pdb_issue_20766[0]>(6)test_function()
    -> print('pdb %d: %s' % (i, sess._previous_sigint_handler))
    (Pdb) continue
    pdb 1: <built-in function default_int_handler>
    > <doctest test.test_pdb.test_pdb_issue_20766[0]>(5)test_function()
    -> sess.set_trace(sys._getframe())
    (Pdb) continue
    pdb 2: <built-in function default_int_handler>
    """

class PdbTestCase(unittest.TestCase):

    def run_pdb(self, script, commands):
        """Run 'script' lines with pdb and the pdb 'commands'."""
        filename = 'main.py'
        with open(filename, 'w') as f:
            f.write(textwrap.dedent(script))
        self.addCleanup(support.unlink, filename)
        self.addCleanup(support.rmtree, '__pycache__')
        cmd = [sys.executable, '-m', 'pdb', filename]
        stdout = stderr = None
        with subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                   stdin=subprocess.PIPE,
                                   stderr=subprocess.STDOUT,
                                   ) as proc:
            stdout, stderr = proc.communicate(str.encode(commands))
        stdout = stdout and bytes.decode(stdout)
        stderr = stderr and bytes.decode(stderr)
        return stdout, stderr

    def _assert_find_function(self, file_content, func_name, expected):
        file_content = textwrap.dedent(file_content)

        with open(support.TESTFN, 'w') as f:
            f.write(file_content)

        expected = None if not expected else (
            expected[0], support.TESTFN, expected[1])
        self.assertEqual(
            expected, pdb.find_function(func_name, support.TESTFN))

    def test_find_function_empty_file(self):
        self._assert_find_function('', 'foo', None)

    def test_find_function_found(self):
        self._assert_find_function(
            """\
            def foo():
                pass

            def bar():
                pass

            def quux():
                pass
            """,
            'bar',
            ('bar', 4),
        )

    def test_issue7964(self):
        # open the file as binary so we can force \r\n newline
        with open(support.TESTFN, 'wb') as f:
            f.write(b'print("testing my pdb")\r\n')
        cmd = [sys.executable, '-m', 'pdb', support.TESTFN]
        proc = subprocess.Popen(cmd,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            )
        self.addCleanup(proc.stdout.close)
        stdout, stderr = proc.communicate(b'quit\n')
        self.assertNotIn(b'SyntaxError', stdout,
                         "Got a syntax error running test script under PDB")

    def test_issue13183(self):
        script = """
            from bar import bar

            def foo():
                bar()

            def nope():
                pass

            def foobar():
                foo()
                nope()

            foobar()
        """
        commands = """
            from bar import bar
            break bar
            continue
            step
            step
            quit
        """
        bar = """
            def bar():
                pass
        """
        with open('bar.py', 'w') as f:
            f.write(textwrap.dedent(bar))
        self.addCleanup(support.unlink, 'bar.py')
        stdout, stderr = self.run_pdb(script, commands)
        self.assertTrue(
            any('main.py(5)foo()->None' in l for l in stdout.splitlines()),
            'Fail to step into the caller after a return')

    def test_issue13210(self):
        # invoking "continue" on a non-main thread triggered an exception
        # inside signal.signal

        # raises SkipTest if python was built without threads
        support.import_module('threading')

        with open(support.TESTFN, 'wb') as f:
            f.write(textwrap.dedent("""
                import threading
                import pdb

                def start_pdb():
                    pdb.Pdb(readrc=False).set_trace()
                    x = 1
                    y = 1

                t = threading.Thread(target=start_pdb)
                t.start()""").encode('ascii'))
        cmd = [sys.executable, '-u', support.TESTFN]
        proc = subprocess.Popen(cmd,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            )
        self.addCleanup(proc.stdout.close)
        stdout, stderr = proc.communicate(b'cont\n')
        self.assertNotIn('Error', stdout.decode(),
                         "Got an error running test script under PDB")

    def test_issue16180(self):
        # A syntax error in the debuggee.
        script = "def f: pass\n"
        commands = ''
        expected = "SyntaxError:"
        stdout, stderr = self.run_pdb(script, commands)
        self.assertIn(expected, stdout,
            '\n\nExpected:\n{}\nGot:\n{}\n'
            'Fail to handle a syntax error in the debuggee.'
            .format(expected, stdout))

    def test_readrc_kwarg(self):
        script = textwrap.dedent("""
            import pdb; pdb.Pdb(readrc=False).set_trace()

            print('hello')
        """)

        save_home = os.environ.pop('HOME', None)
        try:
            with support.temp_cwd():
                with open('.pdbrc', 'w') as f:
                    f.write("invalid\n")

                with open('main.py', 'w') as f:
                    f.write(script)

                cmd = [sys.executable, 'main.py']
                proc = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stdin=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                with proc:
                    stdout, stderr = proc.communicate(b'q\n')
                    self.assertNotIn("NameError: name 'invalid' is not defined",
                                  stdout.decode())

        finally:
            if save_home is not None:
                os.environ['HOME'] = save_home

    def tearDown(self):
        support.unlink(support.TESTFN)

def load_tests(*args):
    from test import test_pdb
    suites = [unittest.makeSuite(PdbTestCase), doctest.DocTestSuite(test_pdb)]
    return unittest.TestSuite(suites)

if __name__ == '__main__':
    unittest.main()


"""
python gefsv12_ensemble_probability_weighted.py cyyyymmddhh clead cthresh
"""

import os, sys
from datetime import datetime
import numpy as np
import numpy.ma as ma
import _pickle as cPickle
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from matplotlib import rcParams
from mpl_toolkits.basemap import Basemap, interp
from mpl_toolkits.axes_grid1 import make_axes_locatable
import _pickle as cPickle
import scipy.stats as stats

rcParams['xtick.labelsize']='medium'
rcParams['ytick.labelsize']='medium'
rcParams['legend.fontsize']='large'

# =====================================================================
# =====================================================================

# ---- inputs from command line

cyyyymmddhh = sys.argv[1] #
clead = sys.argv[2] # 018 not 18
cthresh = sys.argv[3]  #0.254, 1.0,  5.0, 10.0, 25.0
jnd = 1350
ind = 1110

cyear = cyyyymmddhh[0:4]
rthresh = float(cthresh)
cmonth = cyyyymmddhh[4:6]
imonth = int(cmonth)-1
cmonths = ['Jan','Feb','Mar','Apr','May','Jun',\
    'Jul','Aug','Sep','Oct','Nov','Dec']
ccmonth = cmonths[imonth]
cdomain = 'conus'
#master_directory_fullfield_output = '/Volumes/NBM/conus_gefsv12/fullfield/'
master_directory_thinned_output = '/Volumes/NBM/conus_gefsv12/thinned/'
dressed = False
plot_fullfield = True
weighted = True
makeplot = True

# ---- read the stored netCDF full-field probabilities, quantile mapped

#if weighted == True:
#    if dressed == True:
#        infile = master_directory_fullfield_output +\
#            ccmonth+cyear+'_use99_lead'+clead+\
#            '_probabilities_weighted_dressed_fullfield.nc'
#    else:
#        infile = master_directory_fullfield_output +\
#            ccmonth+cyear+'_use99_lead'+clead+\
#            '_probabilities_weighted_fullfield.nc'
#else:
#    infile = master_directory_fullfield_output +\
#        ccmonth+cyear+'_use99_lead'+clead+\
#        '_probabilities_fullfield.nc'

if weighted == True:
    if dressed == True:
        infile = master_directory_thinned_output +\
            ccmonth+cyear+'_use99_lead'+clead+\
            '_probabilities_weighted_dressed_thinned.nc'
    else:
        infile = master_directory_thinned_output +\
            ccmonth+cyear+'_use99_lead'+clead+\
            '_probabilities_weighted_thinned.nc'
else:
    infile = master_directory_thinned_output +\
        ccmonth+cyear+'_use99_lead'+clead+\
        '_probabilities_thinned.nc'

print ('reading from ', infile)
nc = Dataset(infile,'r')
yyyymmddhh = nc.variables['yyyymmddhh_init'][:]
idd = np.where(yyyymmddhh == int(cyyyymmddhh))[0]
print ('yyyymmddhh = ', yyyymmddhh)
thresholds = nc.variables['thresholdv'][:]
if rthresh == 0.254:
    itt = 0
else:
    itt = int(np.where(thresholds == rthresh)[0])
print ('itt = ', itt)
lons_ndfd = nc.variables['lons'][:,:]
lats_ndfd = nc.variables['lats'][:,:]
prob_qmapped_fullfield = \
    np.squeeze(nc.variables['probability_qmapped'][idd,itt,:,:])
nc.close()
ny_ndfd, nx_ndfd = np.shape(lons_ndfd)

# CONUS
#m = Basemap(llcrnrlon=233.7234,llcrnrlat=19.229,\
#    urcrnrlon = 300.95782, urcrnrlat = 54.37279,\
#    projection='lcc',lat_1=25.,lat_2=25.,lon_0=265.,\
#    resolution ='l',area_thresh=1000.)

# Western US
m = Basemap(llcrnrlon=235.,llcrnrlat=30.,\
    urcrnrlon = 260., urcrnrlat = 45.,\
    projection='lcc',lat_1=25.,lat_2=25.,lon_0=265.,\
    resolution ='l',area_thresh=1000.)

# Pac Northwest
#m = Basemap(llcrnrlon=233.7234,llcrnrlat=37.,\
#    urcrnrlon = 260., urcrnrlat = 54.37279,\
#    projection='lcc',lat_1=25.,lat_2=25.,lon_0=265.,\
#    resolution ='l',area_thresh=1000.)

#m = Basemap(llcrnrlon=245.,llcrnrlat=35.,
#    urcrnrlon = 260., urcrnrlat = 43.,\
#    projection='lcc',lat_1=25.,lat_2=25.,lon_0=255.,\
#    resolution ='l',area_thresh=1000.)

# Cali
#m = Basemap(llcrnrlon=235.,llcrnrlat=31.,
#    urcrnrlon = 245., urcrnrlat = 42.,\
#    projection='lcc',lat_1=25.,lat_2=25.,lon_0=245.,\
#    resolution ='l',area_thresh=1000.)

#m = Basemap(llcrnrlon=230.,llcrnrlat=22.,
#    urcrnrlon = 290., urcrnrlat = 51.,\
#    projection='lcc',lat_1=25.,lat_2=25.,lon_0=255.,\
#    resolution ='l',area_thresh=1000.)

clevs = [0.0,0.05,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,0.93,0.97,1.0]
colorst = ['White','#E4FFFF','#C4E8FF','#8FB3FF','#D8F9D8',\
    '#A6ECA6','#42F742','Yellow','Gold','Orange','#FCD5D9','#F6A3AE',\
    '#FA5257','Orchid','#AD8ADB','#A449FF','LightGray']
parallels = np.arange(20.,60.01,5.)
meridians = np.arange(-140.,-49.95,5.)
lons = lons_ndfd
lats = lats_ndfd
prob_qmapped = np.squeeze(prob_qmapped_fullfield)
ctype = ''

if makeplot == True:
    x, y = m(lons, lats)

    # ---- plot the quantile-mapped probability

    #fig = plt.figure(figsize=(9,6.5))
    fig = plt.figure(figsize=(9.,6.5))
    #fig = plt.figure(figsize=(6.,7.7))
    #axloc = [0.06,0.15,0.92,0.8]
    axloc = [0.06,0.11,0.92,0.8]
    ax1 = fig.add_axes(axloc)
    cleadb = str(int(clead)-6)
    #title = ctype+' Quantile-mapped GEFSv12 P(obs > '+\
    #    cthresh+' mm), IC = '+cyyyymmddhh+', lead = '+clead+' h'
    if weighted == True:
        if dressed == True:
            title = 'Weighted, dressed quantile-mapped GEFSv12 P(obs > '+\
                cthresh+' mm),\nIC = '+cyyyymmddhh+', lead = '+clead+' h'
            plot_title = 'qmapped_chist_dressed_probability_'+cthresh+'mm_'+\
                cyyyymmddhh+'_lead'+clead+'.png'
        else:
            title = 'Weighted, quantile-mapped GEFSv12 P(obs > '+\
                cthresh+' mm),\nIC = '+cyyyymmddhh+', lead = '+clead+' h'
            plot_title = 'qmapped_chist_probability_'+cthresh+'mm_'+\
                cyyyymmddhh+'_lead'+clead+'.png'
    else:
        title = 'Quantile-mapped GEFSv12 P(obs > '+\
            cthresh+' mm),\nIC = '+cyyyymmddhh+', lead = '+clead+' h'
        plot_title = 'qmapped_probability_'+cthresh+'mm_'+\
            cyyyymmddhh+'_lead'+clead+'.png'

    ax1.set_title(title, fontsize=14,color='Black')
    CS2 = m.contourf(x, y, prob_qmapped[:,:], clevs,\
        cmap=None, colors=colorst, extend='both')

    m.drawcoastlines(linewidth=0.8,color='Gray')
    m.drawcountries(linewidth=0.8,color='Gray')
    m.drawstates(linewidth=0.8,color='Gray')
    m.drawcounties(linewidth=0.3,color='LightGray')

    #m.drawparallels(parallels,labels=[1,0,0,0],\
    #    color='Gray',fontsize=8,linewidth=0.2)
    #meridians = np.arange(220.,360.,5.)
    #m.drawmeridians(meridians,labels=[0,0,0,1],\
    #    color='Gray',fontsize=8,linewidth=0.2)

    cax = fig.add_axes([0.06,0.06,0.92,0.02])
    cb = plt.colorbar(CS2,orientation='horizontal',cax=cax,\
        drawedges=True,ticks=clevs,format='%g')
    cb.ax.tick_params(labelsize=7)
    cb.set_label('Probability',fontsize=9)

    fig.savefig(plot_title, dpi=300)
    plt.close()
    print ('saving plot to file = ',plot_title)

#


from __future__ import annotations

from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from warnings import warn

import numpy as np
import pandas as pd
from numpy.polynomial.polynomial import polyfit, polyval
from numpy.typing import ArrayLike
from pandas import DataFrame
from scipy.interpolate import UnivariateSpline as USpline
from scipy.optimize import curve_fit
from typing_extensions import Literal

from empyricalRMT._constants import (
    DEFAULT_POLY_DEGREE,
    DEFAULT_SPLINE_DEGREE,
    DEFAULT_SPLINE_DEGREES,
    DEFAULT_SPLINE_SMOOTH,
    DEFAULT_SPLINE_SMOOTHS,
)
from empyricalRMT._types import fArr
from empyricalRMT.detrend import emd_detrend
from empyricalRMT.exponentials import gompertz

SPLINE_DICT = {3: "cubic", 4: "quartic", 5: "quintic"}

class SmoothMethod(Enum):
    Polynomial = "poly"
    Spline = "spline"
    Exponential = "exp"
    Gompertz = "gompertz"
    GOE = "goe"
    Poisson = "poisson"

    @classmethod
    def validate(cls, s: str | SmoothMethod) -> SmoothMethod:
        try:
            if isinstance(s, str):
                return cls(s)
            return s
        except Exception as e:
            values = [e.value for e in cls]
            raise ValueError(f"SmoothMethod must be one of {values}") from e

SmoothArg = Union[List[float], Literal["heuristic"]]

def _spline_name(i: int) -> str:
    return SPLINE_DICT[i] if SPLINE_DICT.get(i) is not None else f"deg{i}"

class Smoother:
    def __init__(self, eigenvalues: ArrayLike):
        """Initialize a Smoother.

        Parameters
        ----------
        eigenvalues: ndarray
            Eigenvalues for fitting to the step function.
        """
        try:
            eigs = np.array(eigenvalues).astype(np.float64)
        except BaseException as e:
            raise ValueError("Could not convert eigenvalues into numpy array.") from e
        if len(eigs) != len(eigs.reshape((-1,))):
            raise ValueError("Input array must be one-dimensional.")
        eigs = eigs.flatten()
        self._eigs = np.sort(eigs)

    def fit(
        self,
        smoother: SmoothMethod = SmoothMethod.Polynomial,
        degree: int = DEFAULT_POLY_DEGREE,
        spline_smooth: float = DEFAULT_SPLINE_SMOOTH,
        detrend: bool = False,
        return_callable: bool = False,
    ) -> Tuple[fArr, fArr, Optional[Callable[[fArr], fArr]]]:
        """Computer the specified smoothing function values for a set of eigenvalues.

        Parameters
        ----------
        eigs: ndarray
            The sorted eigenvalues

        smoother: "poly" | "spline" | "gompertz" | lambda
            The type of smoothing function used to fit the step function

        degree: int
            The degree of the polynomial or spline

        spline_smooth: float
            The smoothing factors passed into scipy.interpolate.UnivariateSpline

        detrend: bool
            Whether or not to perform EMD detrending before returning the
            unfolded eigenvalues.

        return_callable: bool
            If true, return a function that closes over the fit parameters so
            that, e.g., additional values can be fit later.

        Returns
        -------
        unfolded: ndarray
            the unfolded eigenvalues

        steps: ndarray
            the step-function values
        """
        eigs = self._eigs
        # steps = _step_function_fast(eigs, eigs)
        steps = np.arange(0, len(eigs)) + 1
        self.__validate_args(smoother=smoother, degree=degree, spline_smooth=spline_smooth)

        if smoother is SmoothMethod.Polynomial:
            poly_coef = polyfit(eigs, steps, degree)
            unfolded = polyval(eigs, poly_coef)
            func = lambda x: polyval(x, poly_coef) if return_callable else None
            if detrend:
                unfolded = emd_detrend(unfolded)
            return unfolded, steps, func  # type: ignore

        if smoother is SmoothMethod.Spline:
            k = DEFAULT_SPLINE_DEGREE
            try:
                k = int(degree)
            except BaseException as e:
                print(ValueError("Cannot convert spline degree to int."))
                raise e
            if spline_smooth == "heuristic":
                s = len(eigs) * np.var(eigs, ddof=1)
                spline = USpline(eigs, steps, k=k, s=s)
            elif spline_smooth is not None:
                if not isinstance(spline_smooth, float):
                    raise ValueError("Spline smoothing factor must be a float")
                spline = USpline(eigs, steps, k=k, s=spline_smooth)
            else:
                raise ValueError(
                    "Unreachable: All possible spline_smooth arguments should have been handled."
                )
                spline = USpline(eigs, steps, k=k, s=spline_smooth)
            func = spline if return_callable else None
            unfolded = np.array(spline(eigs))
            if detrend:
                unfolded = emd_detrend(unfolded)
            return unfolded, steps, func  # type: ignore

        if smoother is SmoothMethod.Exponential:
            func = lambda t, a, b: a * np.exp(-b * t)  # type: ignore
            [a, b], cov = curve_fit(
                func,
                eigs,
                steps,
                p0=(len(eigs) / 2, 0.5),
            )
            unfolded = func(eigs, a, b)  # type: ignore
            if detrend:
                unfolded = emd_detrend(unfolded)
            return unfolded, steps, func  # type: ignore

        if smoother is SmoothMethod.Gompertz:
            # use steps[end] as guess for the asymptote, a, of gompertz curve
            [a, b, c], cov = curve_fit(gompertz, eigs, steps, p0=(steps[-1], 1, 1))
            func = lambda x: gompertz(x, a, b, c) if return_callable else None
            unfolded = gompertz(eigs, a, b, c)
            if detrend:
                unfolded = emd_detrend(unfolded)
            return unfolded, steps, func  # type: ignore
        raise RuntimeError("Unreachable!")

    def fit_all(
        self,
        poly_degrees: List[int] = [],
        spline_smooths: SmoothArg = [],
        spline_degrees: List[int] = DEFAULT_SPLINE_DEGREES,
        gompertz: bool = False,
        detrend: bool = False,
    ) -> Tuple[DataFrame, DataFrame, DataFrame, Dict[str, Callable]]:
        """unfold eigenvalues for all specified smoothers

        Parameters
        ----------
        poly_degrees: List[int]
            the polynomial degrees for which to compute fits.
            Default [3, 4, 5, 6, 7, 8, 9, 10, 11]

        spline_smooths: List[float] | "heuristic"
            If a list of floats, the smoothing factors, s, passed into
            scipy.interpolate.UnivariateSpline.
            If "heuristic", choose a set of smoothing factors scaled to the length of the
            eigenvalues, that, on GOE eigenvalues, tend to result in a range of fits
            varying from highly flexible (nearly interpolated) to about the flexibility of
            a cubic or quartic. As the number of eigenvalues starts to go below about 300,
            an increasing number of practically-identical, redundant splines will be fit
            with this option, and manual inspection or non-heuristic specification of
            spline smoothing factors is strongly recommended.

        spline_degrees: List[int]
            A list of ints determining the degrees of scipy.interpolate.UnivariateSpline
            fits. Default [3]

        Returns
        -------
        unfoldeds: DataFrame
            DataFrame of unfolded eigenvalues for each set of fit parameters, e.g. where
            each column contains a name indicating the fitting parameters, with the values
            of that column being the (sorted) unfolded eigenvalues.

        sqes: DataFrame
            DataFrame of mean-squared error of fits, where each column contains a name
            indicating the fitting parameters and smoother, with the values of
            the column being the mean of the squared residuals of the fit

        smoother_map: dict
            A dict of {col_name: closure} for accessing the fitted smoothers later.
        """
        # construct dataframes to hold all info
        col_names, unfoldeds, spacings, sqes = [], [], [], []
        smoother_map = {}
        for d in poly_degrees:
            col_name = f"poly_{d}"
            unfolded, steps, closure = self.fit(
                smoother=SmoothMethod.Polynomial, degree=d, return_callable=True, detrend=detrend
            )
            col_names.append(col_name)
            sqes.append(np.mean((unfolded - steps) ** 2))
            unfolded = np.sort(unfolded)  # Important!
            unfoldeds.append(unfolded)
            spacings.append(np.diff(unfolded))
            smoother_map[col_name] = closure

        if spline_smooths == "heuristic":
            for s in DEFAULT_SPLINE_SMOOTHS:
                for d in spline_degrees:
                    col_name = f"{_spline_name(d)}-spline_" "{:1.2f}_heuristic".format(s)
                    unfolded, steps, closure = self.fit(
                        smoother=SmoothMethod.Spline,
                        spline_smooth=len(self._eigs) ** s,
                        degree=d,
                        return_callable=True,
                        detrend=detrend,
                    )
                    col_names.append(col_name)
                    sqes.append(np.mean((unfolded - steps) ** 2))
                    unfolded = np.sort(unfolded)
                    unfoldeds.append(unfolded)
                    spacings.append(np.diff(unfolded))
                    smoother_map[col_name] = closure
        else:
            for s in spline_smooths:  # type: ignore
                for d in spline_degrees:
                    col_name = f"{_spline_name(d)}-spline_" "{:1.3f}".format(s)
                    unfolded, steps, closure = self.fit(
                        smoother=SmoothMethod.Spline,
                        spline_smooth=s,
                        degree=d,
                        return_callable=True,
                        detrend=detrend,
                    )
                    col_names.append(col_name)
                    sqes.append(np.mean((unfolded - steps) ** 2))
                    unfolded = np.sort(unfolded)
                    unfoldeds.append(unfolded)
                    spacings.append(np.diff(unfolded))
                    smoother_map[col_name] = closure
        if gompertz:
            unfolded, steps, closure = self.fit(
                smoother=SmoothMethod.Gompertz, return_callable=True, detrend=detrend
            )
            col_names.append("gompertz")
            sqes.append(np.mean((unfolded - steps) ** 2))
            unfolded = np.sort(unfolded)
            unfoldeds.append(unfolded)
            spacings.append(np.diff(unfolded))
            smoother_map["gompertz"] = closure
        unfoldeds = pd.DataFrame(data=unfoldeds, index=col_names).T
        spacings = pd.DataFrame(data=spacings, index=col_names).T
        sqes = pd.DataFrame(data=sqes, index=col_names).T
        return unfoldeds, spacings, sqes, smoother_map  # type: ignore

    @staticmethod
    def _get_smoother_names(
        poly_degrees: List[int],
        spline_smooths: SmoothArg,
        spline_degrees: List[int] = [3],
        gompertz: bool = True,
    ) -> List[str]:
        """If arguments are arrays, generate names (unique identifiers) for each smoother
        + smoother parameters. Otherwise, just return the name for indexing into the report.
        """

        col_names = []
        if isinstance(poly_degrees, list):
            for d in poly_degrees:
                col_names.append(f"poly_{d}")
        else:
            raise ValueError("poly_degrees must be a list of int values")

        if spline_smooths == "heuristic":
            for s in DEFAULT_SPLINE_SMOOTHS:
                if not isinstance(spline_degrees, list):
                    raise ValueError("spline_degrees must be a list of integer values")
                for deg in spline_degrees:
                    col_name = f"{_spline_name(deg)}-spline_" "{:1.3f}_heuristic".format(s)
                    col_names.append(col_name)
        else:
            try:
                spline_smooths = list(spline_smooths)  # type: ignore
            except Exception as e:
                raise ValueError(f"Error converting `spline_smooths` to list: {e}")
            if isinstance(spline_smooths, list):
                for s in spline_smooths:
                    if not isinstance(spline_degrees, list):
                        raise ValueError("spline_degrees must be a list of integer values")
                    for deg in spline_degrees:
                        col_name = f"{_spline_name(deg)}-spline_" "{:1.3f}".format(s)
                        col_names.append(col_name)
            else:
                raise ValueError("spline_smooths must be a list of float values")

        if gompertz is True:
            col_names.append("gompertz")
        return col_names

    def __validate_args(self, **kwargs: Any) -> None:
        """throw an error if smoother args are in any way invalid"""
        smoother = SmoothMethod.validate(kwargs.get("smoother"))  # type: ignore
        degree = kwargs.get("degree")
        spline_smooth = kwargs.get("spline_smooth")
        emd = kwargs.get("detrend")  # TODO: implement
        method = kwargs.get("method")

        if smoother is SmoothMethod.Polynomial:
            if degree is None:
                warn(
                    "No degree set for polynomial unfolding."
                    f"Will default to polynomial of degree {DEFAULT_POLY_DEGREE}.",
                    category=UserWarning,
                )
            if not isinstance(degree, int):
                raise ValueError("Polynomial degree must be of type `int`")
            if degree < 3:
                raise ValueError("Unfolding polynomial must have minimum degree 3.")
        elif smoother is SmoothMethod.Spline:
            spline_degree = degree
            if degree is None:
                warn(
                    "No degree set for spline unfolding. Will default to "
                    f"spline of degree {DEFAULT_SPLINE_DEGREE}.",
                    category=UserWarning,
                )
            if not isinstance(spline_degree, int) or spline_degree > 5:
                raise ValueError("Degree of spline must be an int <= 5")
            if spline_smooth is not None and spline_smooth != "heuristic":
                spline_smooth = float(spline_smooth)
        elif smoother in [
            SmoothMethod.Gompertz,
            SmoothMethod.Exponential,
            SmoothMethod.GOE,
            SmoothMethod.Poisson,
        ]:
            pass  # just allow this for now
        elif callable(smoother):
            # NOTE: above is not a great check, but probably good enough for our purposes
            # https://stackoverflow.com/questions/624926/how-do-i-detect-whether-a-python-variable-is-a-function#comment437753_624939
            raise NotImplementedError("Custom fit functions not currently implemented.")
        else:
            raise ValueError("Unrecognized smoother argument.")

        if emd is not None and not isinstance(emd, bool):
            raise ValueError("`detrend` can be only a boolean or undefined (None).")

        if method is None or method == "auto" or method == "manual":
            pass
        else:
            raise ValueError("`method` must be one of 'auto', 'manual', or 'None'")


import fnmatch
import functools
import json
import logging
import re
import struct
import time

import numpy as np
import pynq.lib.dma
from ament_index_python.packages import get_package_share_directory
from pynq import MMIO, Overlay, allocate
from pynq.pl_server import hwh_parser

class FpgaDriver:
    def __init__(self, user_ip):
        self.logger = logging.getLogger("PYNQ-Driver")
        logger_handler = logging.StreamHandler()
        logger_handler.setFormatter(
            logging.Formatter("[%(levelname)s] [%(name)s]: %(message)s")
        )
        self.logger.addHandler(logger_handler)

        # FPGA bit file
        self.bit_file = "/home/ubuntu/dev_ws/traffic_light_recognition.bit"
        self.ovl = Overlay(self.bit_file, download=False)
        self.ovl.is_loaded = self.is_loaded
        # User IP
        self.user_ip = user_ip

    def is_loaded(self):
        return True

    def update_map(self, user_ip):
        if not user_ip:
            return
        io_maps_json = open(
            get_package_share_directory("traffic_light_recognition_fpga_node") + "/io_maps.json",
            "r",
        )
        io_maps = json.load(io_maps_json)
        io_maps_json.close()
        self.map_num_str = str(io_maps["map_num"][user_ip])

        # Define the arrays that hold the information about the input and output
        # signals. If the value corresponds to a double or float variable, it is
        # in IEEE-754 representation
        self.in_map = io_maps["maps"][self.map_num_str]["input"]
        self.out_map = io_maps["maps"][self.map_num_str]["output"]
        self.has_axis_in = self.has_axis(self.in_map)
        self.has_axis_out = self.has_axis(self.out_map)
        if self.has_axis_in or self.has_axis_out:
            self.dma_name = self.find_dma_ip_name()

    def find_dma_ip_name(self):
        hwh_file = hwh_parser.get_hwh_name(self.bit_file)
        hwh = hwh_parser.HWH(hwh_file)
        pins = fnmatch.filter(hwh.pins.keys(), "{}/*TDATA".format(self.user_ip))
        for pin in pins:
            dma_match = re.search(r"axi_dma_\d+", hwh.pins[pin])
            if dma_match is not None:
                return dma_match.group()

    @property
    def user_ip(self):
        return self.__user_ip

    @user_ip.setter
    def user_ip(self, ip_name):
        self.__user_ip = ip_name
        self.update_map(ip_name)

    def program_fpga(self):
        # Program the FPGA and set the Overlay
        print(f"Downloading bit file {self.bit_file}...")
        self.ovl.download()
        print("Bit file downloaded")

    def setup_fpga(self):
        # Set the MMIO and DMA classes
        regs = self.ovl.ip_dict[self.user_ip]
        phys_addr = regs["phys_addr"]
        addr_range = regs["addr_range"]
        print("User IP phys_addr = 0x{:X}".format(phys_addr))
        print("User IP addr_range = {}".format(addr_range))
        self.mmio = MMIO(phys_addr, addr_range)
        if self.has_axis_in or self.has_axis_out:
            self.dma = getattr(self.ovl, self.dma_name)
        if self.has_axis_in:
            self.input_buffer = self.allocate_dma_buffer(self.in_map)
        if self.has_axis_out:
            self.output_buffer = self.allocate_dma_buffer(self.out_map)

    def has_axis(self, io_map):
        has_axis = False
        for signal_name in io_map.keys():
            has_axis |= io_map[signal_name]["protocol"] == "stream"
        return has_axis

    def allocate_dma_buffer(self, io_map):
        signal_type = ""
        n_bits = 0
        arr_size = 0
        for signal_name in io_map.keys():
            if io_map[signal_name]["protocol"] == "stream":
                signal_type = io_map[signal_name]["type"]
                n_bits = io_map[signal_name]["n_bits"]
                arr_size = io_map[signal_name]["n_elem"]
                break

        # case 1 - int32
        if signal_type == "int" and n_bits == 32:
            allocated_buffer = allocate(shape=(arr_size,), dtype=np.int32)
        # case 2 - uint8
        if signal_type == "int" and n_bits == 8:
            allocated_buffer = allocate(shape=(arr_size,), dtype=np.uint8)
        # case 3 - uint64
        if signal_type == "int" and n_bits == 64:
            allocated_buffer = allocate(shape=(arr_size,), dtype=np.uint64)
        # case 4 - float32
        if signal_type == "float" and n_bits == 32:
            allocated_buffer = allocate(shape=(arr_size,), dtype=np.float32)
        return allocated_buffer

    def process_input(self, signal_name, msg):
        # Receive the value from the msg and send it to the FPGA
        # Need to figure out if the input is an array, and the input type
        signal_type = self.in_map[signal_name]["type"]
        is_arr = self.in_map[signal_name]["arr"]
        n_bits = self.in_map[signal_name]["n_bits"]
        protocol = self.in_map[signal_name]["protocol"]

        if protocol == "lite":

            # case 1 - Signal represents an array
            if is_arr:
                arr_head = self.in_map[signal_name]["addr"]
                arr_size = self.in_map[signal_name]["n_elem"]
                if signal_type == "int":
                    if n_bits == 32:  # int32
                        # 1 array element per address
                        for i in range(arr_size):
                            self.mmio.write(arr_head + 4 * i, int(msg[i]))
                    elif n_bits == 8:  # uint8
                        # 4 array elements per address
                        byte_arr = bytes(msg)
                        lower_lim = 0
                        upper_lim = 4
                        n_exact = arr_size // 4
                        n_last = arr_size % 4
                        padding = 4 - n_last
                        addr = arr_head
                        for i in range(n_exact):
                            partition = byte_arr[lower_lim:upper_lim]
                            self.mmio.write(addr, partition)
                            lower_lim += 4
                            upper_lim += 4
                            addr += 4
                        if n_last != 0:
                            partition = byte_arr[arr_size - n_last : arr_size]
                            pad = bytes(padding)
                            partition += pad
                            self.mmio.write(addr, partition)
                elif signal_type == "float":
                    if n_bits == 32:  # Float
                        # 1 array element per address
                        for i in range(arr_size):
                            ieee_rep = self.dec_to_ieee(msg[i], "float")
                            self.mmio.write(arr_head + 4 * i, ieee_rep)
                    elif n_bits == 64:  # Double
                        # 1 array element per 2 addresses
                        for i in range(arr_size):
                            ieee_rep_hi, ieee_rep_lo = self.dec_to_ieee(msg[i], "double")
                            self.mmio.write(arr_head + 8 * i, ieee_rep_lo)
                            self.mmio.write(arr_head + 8 * i + 4, ieee_rep_hi)

            # case 2 - Signal does not represent an array
            else:
                addr = self.in_map[signal_name]["addr"]
                if signal_type == "int":
                    self.mmio.write(addr, int(msg))
                elif signal_type == "float":
                    if n_bits == 32:  # Float
                        ieee_rep = self.dec_to_ieee(msg, "float")
                        self.mmio.write(addr, ieee_rep)
                    elif n_bits == 64:  # Double
                        ieee_rep_hi, ieee_rep_lo = self.dec_to_ieee(msg, "double")
                        self.mmio.write(addr, ieee_rep_lo)
                        self.mmio.write(addr + 4, ieee_rep_hi)

        elif protocol == "stream":
            self.input_buffer[:] = msg[:]

    def process_output(self, signal_name):
        # Read the value from the FPGA and return it in a format that can be assigned to a msg
        # Need to figure out if the output is an array, and the output type
        signal_type = self.out_map[signal_name]["type"]
        is_arr = self.out_map[signal_name]["arr"]
        n_bits = self.out_map[signal_name]["n_bits"]
        protocol = self.out_map[signal_name]["protocol"]

        if protocol == "lite":

            if is_arr:
                arr_head = self.out_map[signal_name]["addr"]
                arr_size = self.out_map[signal_name]["n_elem"]
                val = [0] * arr_size
                if signal_type == "int":
                    if n_bits == 32:  # int32
                        for i in range(arr_size):
                            val[i] = self.mmio.read(arr_head + 4 * i)
                        return val
                    elif n_bits == 8:  # uint8
                        # 4 array elements per address
                        n_exact = arr_size // 4
                        n_last = arr_size % 4
                        padding = 4 - n_last
                        addr = arr_head
                        for i in range(n_exact):
                            result = self.mmio.read(addr).to_bytes(4, "little")
                            for j in range(4):
                                val[j + i] = int(result[j])
                            addr += 4
                        if n_last != 0:
                            result = self.mmio.read(addr).to_bytes(n_last, "little")
                            for j in range(n_last):
                                val[arr_size - n_last + j] = int(result[j])
                        return val
                elif signal_type == "float":
                    if n_bits == 32:  # Float
                        for i in range(arr_size):
                            val_ieee = self.mmio.read(arr_head + 4 * i)
                            val_ieee = hex(val)[2:]
                            val[i] = self.ieee_to_dec(val_ieee, "float")
                        return val
                    elif n_bits == 64:  # Double
                        for i in range(arr_size):
                            val_ieee_lo = self.mmio.read(arr_head + 8 * i)
                            val_ieee_hi = self.mmio.read(arr_head + 8 * i + 4)
                            val_ieee = hex((val_ieee_hi << 32) | val_ieee_lo)[2:]
                            val[i] = self.ieee_to_dec(val_ieee, "double")
                        return val

        elif protocol == "stream":
            data_size_out = self.out_map[signal_name]["n_elem"]
            if self.out_map[signal_name]["n_bits"] == 64:
                val = np.zeros((data_size_out,), dtype=np.uint64)
            elif signal_type == "float":
                val = np.zeros((data_size_out,), dtype=np.float32)
            else:
                val = [0] * data_size_out
            val[:] = self.output_buffer[:]
            return val

    def setup_dma_buffers(self):
        # Setup the input and output AXI_Stream buffers for DMA transfer

        if self.has_axis_in:
            self.dma.sendchannel.transfer(self.input_buffer)
        if self.has_axis_out:
            self.dma.recvchannel.transfer(self.output_buffer)

        return

    def do_calc(self):
        # Send the AP_START flag
        self.mmio.write(0x00, 1)
        # Wait for streams, if any
        if self.has_axis_in and self.dma.sendchannel.running:
            self.dma.sendchannel.wait()
        if self.has_axis_out and self.dma.recvchannel.running:
            self.dma.recvchannel.wait()

        # Wait for the computation to be finished (AP_DONE)
        self.wait_result()

    def wait_result(self):
        # Wait for AP_DONE
        while not ((self.mmio.read(0x00) >> 1) & 0x0001):
            continue

    def dec_to_ieee(self, num, data_type):
        # Converts a decimal number to IEEE-754 representation
        if data_type == "float":
            num_ieee_hex = struct.pack(">f", num).hex()
            num_ieee = int(num_ieee_hex, 16)
            return num_ieee
        elif data_type == "double":
            num_ieee = struct.pack(">d", num).hex()
            num_hi = int(num_ieee[:8], 16)
            num_lo = int(num_ieee[8:], 16)
            return num_hi, num_lo

    def ieee_to_dec(self, num_ieee, data_type):
        # Converts a number in IEEE-754 representation to its decimal form
        if num_ieee != "0":
            num_in_bytes = bytes.fromhex(num_ieee)
            if data_type == "float":
                num_dec = struct.unpack(">f", num_in_bytes)[0]
            elif data_type == "double":
                num_dec = struct.unpack(">d", num_in_bytes)[0]
        else:
            num_dec = 0
        return num_dec