import torch
import numpy as np

def train_perm_orth(train_loader, model, optimizer, scheduler, criterion, regularizer=None, rho=1E-4, delta=0.5,
                    nu=1E-2, eps=1E-3, tau=1E-2, lagrange_pen=1E-2, perm_flag=True, t_step=40):
    if perm_flag:
        tau_min = 1E-24
        tau_max = 1E-1
        c = None
        lam_lm = []
        for p in optimizer.param_groups[0]['params']:
            lam_lm.append(torch.zeros_like(p))

        k_iter = 0
        ts = torch.empty(len(train_loader), device=model.device).uniform_(0.0, 1.0)
        with torch.no_grad():
            for p in optimizer.param_groups[0]['params']:
                p.data = torch.rand_like(p.data)
                p.data, _, _ = torch.svd(p.data)
        input_cml = []
        target_cml = []
        t_cml = []
        inner_iter = 0
        loss = 0.0
        loss_obj = 0.0
        for iter, (input, target) in enumerate(train_loader):
            t = ts[iter]
            input = input.to(model.device, non_blocking=False)
            target = target.to(model.device, non_blocking=False)
            output = model(input, perm_train=True, t=t)
            input_all = input
            target_all = target

            new_loss = criterion(output, target_all)
            loss_obj += new_loss

            #  This part is for the augmented Lagrangian method
            int_pen = integer_penalty(optimizer.param_groups[0]['params'], lam_lm, lagrange_pen)

            loss += new_loss + int_pen

            inner_iter += 1
            input_cml.append(input.clone())
            target_cml.append(target.clone())
            t_cml.append(t.clone())
            if inner_iter % t_step == 0:
                optimizer.zero_grad()
                loss.backward()
                grad_norm = 0.0
                violator = 0.0
                for p in optimizer.param_groups[0]['params']:
                    param_norm = p.grad.data.norm(2)
                    grad_norm += param_norm.item() ** 2
                    violator += torch.sum((torch.matmul(p.data.t(), p.data) - torch.eye(p.data.shape[0],
                                                                                device=p.device)) ** 2)
                grad_norm = grad_norm ** (1. / 2)

                if c is None:
                    c = loss.clone().item()
                    q_opt = 1
                loss_inner = loss.clone()
                print('Iteration: %03d, Loss %.2E, Objective %.2E, Negative Penalty: %.2E,'
                      'Grad Norm: %.2E, Ortho Violation: %.2E, tau: %.2E' %
                      (k_iter, loss_inner.item(), loss_obj.item(), int_pen.item(), grad_norm, violator.item(), tau))
                #  Compute F for defining Y function
                F_list = []
                with torch.no_grad():
                    for p in optimizer.param_groups[0]['params']:
                        f = torch.matmul(p.grad.data, p.t().data) - torch.matmul(p.data, p.grad.t().data)
                        F_list.append(f)
                    #  Store old parameters
                    params_old = [None] * len(optimizer.param_groups[0]['params'])
                    for idx, param in enumerate(optimizer.param_groups[0]['params']):
                        params_old[idx] = param.clone()
                    grads_old = [p.grad.data.clone() for p in optimizer.param_groups[0]['params']]

                    #  Compute the values of Y(tau) and Y'(tau), store them into the model
                    Y_t, Y_ft_prime = compute_ytau(tau, F_list, optimizer.param_groups[0]['params'])
                    for p, y_t in zip(optimizer.param_groups[0]['params'], Y_t):
                        p.data = y_t.clone()
                loss_inner = 0.0
                for t_2, input_2, target_2 in zip(t_cml, input_cml, target_cml):
                    output = model(input_2, perm_train=True, t=t_2)

                    loss_inner += criterion(output, target_2)
                    int_pen = integer_penalty(optimizer.param_groups[0]['params'], lam_lm, lagrange_pen)
                    loss_inner += int_pen

                optimizer.zero_grad()
                loss_inner.backward()
                grads_new = [p.grad.data.clone() for p in optimizer.param_groups[0]['params']]

                with torch.no_grad():
                    dF_dt = 0.0
                    for g_new, y_ft_p in zip(grads_new, Y_ft_prime):
                        df = g_new * (y_ft_p / torch.norm(y_ft_p.data))
                        df = torch.sum(df)
                        dF_dt += df.item()

                threshold_flag = True
                k_inner = 0
                while threshold_flag:
                    with torch.no_grad():
                        threshold = c + rho * tau * dF_dt
                    if loss_inner.item() >= threshold:
                        # Compute Y for smaller value of tau
                        with torch.no_grad():
                            tau *= delta
                            Y_t, Y_ft_prime = compute_ytau(tau, F_list, optimizer.param_groups[0]['params'])
                            for p, y_t in zip(optimizer.param_groups[0]['params'], Y_t):
                                p.data = y_t.clone()

                        loss_old = loss_inner.clone()
                        loss_inner = 0.0
                        for t_2, input_2, target_2 in zip(t_cml, input_cml, target_cml):
                            output = model(input_2, perm_train=True, t=t_2)
                            loss_inner += criterion(output, target_2)
                            int_pen = integer_penalty(optimizer.param_groups[0]['params'], lam_lm, lagrange_pen)
                            loss_inner += int_pen

                        optimizer.zero_grad()
                        loss_inner.backward()
                        grads_new = [p.grad.data.clone() for p in optimizer.param_groups[0]['params']]

                        k_inner += 1
                        if (loss_inner.item() - loss_old.item()) / (1 + loss_old.item()) < 1E-5:
                            threshold_flag = False
                    else:
                        threshold_flag = False

                with torch.no_grad():
                    c = (nu * q_opt * c + loss_inner.item())
                    q_opt = nu * q_opt + 1
                    c = c / q_opt

                    bb_num = 0.0
                    bb_denom = 0.0
                    yy_sum = 0.0
                    for p_old, g_old, p_new, g_new in zip(params_old, grads_old, optimizer.param_groups[0]['params'],
                                                          grads_new):
                        s_bb = p_new - p_old
                        y_bb = g_new - g_old
                        bb_num += torch.sum(s_bb ** 2)
                        bb_denom += torch.sum(s_bb * y_bb)
                        yy_sum += torch.sum(y_bb ** 2)
                    tau_bb = bb_num / torch.abs(bb_denom)
                    tau_bb = tau_bb.item()
                    tau_bb2 = torch.abs(bb_denom) / yy_sum
                    tau_bb2 = tau_bb2.item()
                    tau_bb = np.minimum(tau_bb, tau_bb2)
                    tau = np.minimum(tau_bb, tau_max)
                    tau = np.maximum(tau, tau_min)
                    lam_lm, lagrange_pen = integer_penalty_update(optimizer.param_groups[0]['params'], lam_lm,
                                                                  lagrange_pen)

                loss_inner = 0.0
                for t_2, input_2, target_2 in zip(t_cml, input_cml, target_cml):
                    output = model(input_2, perm_train=True, t=t_2)
                    loss_obj = criterion(output, target_2)
                    int_pen = integer_penalty(optimizer.param_groups[0]['params'], lam_lm, lagrange_pen)
                    loss_inner += loss_obj + int_pen

                optimizer.zero_grad()
                loss_inner.backward()
                grads_new = [p.grad.data.clone() for p in optimizer.param_groups[0]['params']]

                grad_norm = 0.0
                for g_new in grads_new:
                    gn = g_new.norm(2)
                    grad_norm += gn.item() ** 2
                grad_norm = grad_norm ** (1. / 2)

                k_iter += 1
                input_cml = []
                target_cml = []
                t_cml = []
                loss = 0.0
                loss_obj = 0.0

    model.train()
    loss_sum = 0.0
    correct = 0.0

    change_P = np.nan

    params_before = [None] * len(optimizer.param_groups[0]['params'])
    if nu is not None:
        for idx, param in enumerate(optimizer.param_groups[0]['params']):
            params_before[idx] = param.clone().detach()

        optimizer.step()

        lr = scheduler.get_lr()[0]
        with torch.no_grad():
            for param, param_o in zip(optimizer.param_groups[0]['params'], params_old):
                param.data = 1 / (1 + lr / nu) * (param + lr / nu * param_o)

            output = model(input_all, perm_train=True)
            loss = criterion(output, target_all)
            if regularizer is not None:
                loss += regularizer(model)
            loss_sum += loss.item() * input.size(0)

        pred = output.data.argmax(1, keepdim=True)
        correct += pred.eq(target_all.data.view_as(pred)).sum().item()

    return {
        'loss': loss_sum / len(train_loader.dataset),
        'accuracy': correct * 100.0 / len(train_loader.dataset),
        'change_perm': change_P
    }

def hard_int_penalty(p_list, pen=1E1):
    pen_loss = 0.0
    for p in p_list:
        p_mask = p.data * (p.data <= 0)
        pen_loss += pen * torch.sum(p_mask ** 2)
    return pen_loss

def integer_penalty(p_list, lam_list, mu):
    pen_loss = 0.0
    for p, lam in zip(p_list, lam_list):
        mask = (p - lam / mu) <= 0
        mask_alt = (p - lam / mu) > 0
        p_l = torch.sum((- lam * p + 0.5 * mu * (p ** 2)) * mask)
        p_l += torch.sum((-1/(2 * mu) * lam ** 2) * mask_alt)
        pen_loss += p_l
    return pen_loss

def integer_penalty_update(p_list, lam_list, mu):
    new_lam_list = []
    with torch.no_grad():
        for p, lam in zip(p_list, lam_list):
            upd = lam - mu * p
            new_lam_list.append(upd * (upd > 0))
        new_mu = mu * 1.01
    return new_lam_list, new_mu

def compute_ytau(tau, f_list, p_list):
    y_tau = []
    y_tau_prime = []
    for p, f in zip(p_list, f_list):
        eye = torch.eye(f.shape[0], device=f.device)
        qmat_inv = torch.inverse(eye + tau / 2 * f)
        y_ft = torch.matmul(qmat_inv, eye - tau / 2 * f)
        y_ft = torch.matmul(y_ft, p)
        y_ft_prime = - torch.matmul(qmat_inv, f)
        y_ft_prime = torch.matmul(y_ft_prime, (p + y_ft) / 2)

        y_tau.append(y_ft.clone())
        y_tau_prime.append(y_ft_prime.clone())
    return y_tau, y_tau_prime


"""The test for the History Statistics sensor platform."""
# pylint: disable=protected-access
from datetime import timedelta
import unittest
from unittest.mock import patch

from homeassistant.const import STATE_UNKNOWN
from homeassistant.setup import setup_component
from homeassistant.components.sensor.history_stats import HistoryStatsSensor
import homeassistant.core as ha
from homeassistant.helpers.template import Template
import homeassistant.util.dt as dt_util

from tests.common import init_recorder_component, get_test_home_assistant

class TestHistoryStatsSensor(unittest.TestCase):
    """Test the History Statistics sensor."""

    def setUp(self):
        """Set up things to be run when tests are started."""
        self.hass = get_test_home_assistant()

    def tearDown(self):
        """Stop everything that was started."""
        self.hass.stop()

    def test_setup(self):
        """Test the history statistics sensor setup."""
        self.init_recorder()
        config = {
            'history': {
            },
            'sensor': {
                'platform': 'history_stats',
                'entity_id': 'binary_sensor.test_id',
                'state': 'on',
                'start': '{{ now().replace(hour=0)'
                         '.replace(minute=0).replace(second=0) }}',
                'duration': '02:00',
                'name': 'Test',
            }
        }

        self.assertTrue(setup_component(self.hass, 'sensor', config))

        state = self.hass.states.get('sensor.test')
        self.assertEqual(state.state, STATE_UNKNOWN)

    def test_period_parsing(self):
        """Test the conversion from templates to period."""
        today = Template('{{ now().replace(hour=0).replace(minute=0)'
                         '.replace(second=0) }}', self.hass)
        duration = timedelta(hours=2, minutes=1)

        sensor1 = HistoryStatsSensor(
            self.hass, 'test', 'on', today, None, duration, 'time', 'test')
        sensor2 = HistoryStatsSensor(
            self.hass, 'test', 'on', None, today, duration, 'time', 'test')

        sensor1.update_period()
        sensor1_start, sensor1_end = sensor1._period
        sensor2.update_period()
        sensor2_start, sensor2_end = sensor2._period

        # Start = 00:00:00
        self.assertEqual(sensor1_start.hour, 0)
        self.assertEqual(sensor1_start.minute, 0)
        self.assertEqual(sensor1_start.second, 0)

        # End = 02:01:00
        self.assertEqual(sensor1_end.hour, 2)
        self.assertEqual(sensor1_end.minute, 1)
        self.assertEqual(sensor1_end.second, 0)

        # Start = 21:59:00
        self.assertEqual(sensor2_start.hour, 21)
        self.assertEqual(sensor2_start.minute, 59)
        self.assertEqual(sensor2_start.second, 0)

        # End = 00:00:00
        self.assertEqual(sensor2_end.hour, 0)
        self.assertEqual(sensor2_end.minute, 0)
        self.assertEqual(sensor2_end.second, 0)

    def test_measure(self):
        """Test the history statistics sensor measure."""
        t0 = dt_util.utcnow() - timedelta(minutes=40)
        t1 = t0 + timedelta(minutes=20)
        t2 = dt_util.utcnow() - timedelta(minutes=10)

        # Start     t0        t1        t2        End
        # |--20min--|--20min--|--10min--|--10min--|
        # |---off---|---on----|---off---|---on----|

        fake_states = {
            'binary_sensor.test_id': [
                ha.State('binary_sensor.test_id', 'on', last_changed=t0),
                ha.State('binary_sensor.test_id', 'off', last_changed=t1),
                ha.State('binary_sensor.test_id', 'on', last_changed=t2),
            ]
        }

        start = Template('{{ as_timestamp(now()) - 3600 }}', self.hass)
        end = Template('{{ now() }}', self.hass)

        sensor1 = HistoryStatsSensor(
            self.hass, 'binary_sensor.test_id', 'on', start, end, None,
            'time', 'Test')

        sensor2 = HistoryStatsSensor(
            self.hass, 'unknown.id', 'on', start, end, None, 'time', 'Test')

        sensor3 = HistoryStatsSensor(
            self.hass, 'binary_sensor.test_id', 'on', start, end, None,
            'count', 'test')

        sensor4 = HistoryStatsSensor(
            self.hass, 'binary_sensor.test_id', 'on', start, end, None,
            'ratio', 'test')

        self.assertEqual(sensor1._type, 'time')
        self.assertEqual(sensor3._type, 'count')
        self.assertEqual(sensor4._type, 'ratio')

        with patch('homeassistant.components.history.'
                   'state_changes_during_period', return_value=fake_states):
            with patch('homeassistant.components.history.get_state',
                       return_value=None):
                sensor1.update()
                sensor2.update()
                sensor3.update()
                sensor4.update()

        self.assertEqual(sensor1.state, 0.5)
        self.assertEqual(sensor2.state, None)
        self.assertEqual(sensor3.state, 2)
        self.assertEqual(sensor4.state, 50)

    def test_wrong_date(self):
        """Test when start or end value is not a timestamp or a date."""
        good = Template('{{ now() }}', self.hass)
        bad = Template('{{ TEST }}', self.hass)

        sensor1 = HistoryStatsSensor(
            self.hass, 'test', 'on', good, bad, None, 'time', 'Test')
        sensor2 = HistoryStatsSensor(
            self.hass, 'test', 'on', bad, good, None, 'time', 'Test')

        before_update1 = sensor1._period
        before_update2 = sensor2._period

        sensor1.update_period()
        sensor2.update_period()

        self.assertEqual(before_update1, sensor1._period)
        self.assertEqual(before_update2, sensor2._period)

    def test_wrong_duration(self):
        """Test when duration value is not a timedelta."""
        self.init_recorder()
        config = {
            'history': {
            },
            'sensor': {
                'platform': 'history_stats',
                'entity_id': 'binary_sensor.test_id',
                'name': 'Test',
                'state': 'on',
                'start': '{{ now() }}',
                'duration': 'TEST',
            }
        }

        setup_component(self.hass, 'sensor', config)
        self.assertEqual(self.hass.states.get('sensor.test'), None)
        self.assertRaises(TypeError,
                          setup_component(self.hass, 'sensor', config))

    def test_bad_template(self):
        """Test Exception when the template cannot be parsed."""
        bad = Template('{{ x - 12 }}', self.hass)  # x is undefined
        duration = '01:00'

        sensor1 = HistoryStatsSensor(
            self.hass, 'test', 'on', bad, None, duration, 'time', 'Test')
        sensor2 = HistoryStatsSensor(
            self.hass, 'test', 'on', None, bad, duration, 'time', 'Test')

        before_update1 = sensor1._period
        before_update2 = sensor2._period

        sensor1.update_period()
        sensor2.update_period()

        self.assertEqual(before_update1, sensor1._period)
        self.assertEqual(before_update2, sensor2._period)

    def test_not_enough_arguments(self):
        """Test config when not enough arguments provided."""
        self.init_recorder()
        config = {
            'history': {
            },
            'sensor': {
                'platform': 'history_stats',
                'entity_id': 'binary_sensor.test_id',
                'name': 'Test',
                'state': 'on',
                'start': '{{ now() }}',
            }
        }

        setup_component(self.hass, 'sensor', config)
        self.assertEqual(self.hass.states.get('sensor.test'), None)
        self.assertRaises(TypeError,
                          setup_component(self.hass, 'sensor', config))

    def test_too_many_arguments(self):
        """Test config when too many arguments provided."""
        self.init_recorder()
        config = {
            'history': {
            },
            'sensor': {
                'platform': 'history_stats',
                'entity_id': 'binary_sensor.test_id',
                'name': 'Test',
                'state': 'on',
                'start': '{{ as_timestamp(now()) - 3600 }}',
                'end': '{{ now() }}',
                'duration': '01:00',
            }
        }

        setup_component(self.hass, 'sensor', config)
        self.assertEqual(self.hass.states.get('sensor.test'), None)
        self.assertRaises(TypeError,
                          setup_component(self.hass, 'sensor', config))

    def init_recorder(self):
        """Initialize the recorder."""
        init_recorder_component(self.hass)
        self.hass.start()


"""
This module implements some special functions that commonly appear in
combinatorial contexts (e.g. in power series); in particular,
sequences of rational numbers such as Bernoulli and Fibonacci numbers.

Factorials, binomial coefficients and related functions are located in
the separate 'factorials' module.
"""

from __future__ import print_function, division

from sympy.core import S, Symbol, Rational, Integer, Add, Dummy
from sympy.core.compatibility import as_int, SYMPY_INTS, range
from sympy.core.cache import cacheit
from sympy.core.function import Function, expand_mul
from sympy.core.numbers import E, pi
from sympy.core.relational import LessThan, StrictGreaterThan
from sympy.functions.combinatorial.factorials import binomial, factorial
from sympy.functions.elementary.exponential import log
from sympy.functions.elementary.integers import floor
from sympy.functions.elementary.trigonometric import sin, cos, cot
from sympy.functions.elementary.miscellaneous import sqrt
from sympy.utilities.memoization import recurrence_memo

from mpmath import bernfrac, workprec
from mpmath.libmp import ifib as _ifib

def _product(a, b):
    p = 1
    for k in range(a, b + 1):
        p *= k
    return p

# Dummy symbol used for computing polynomial sequences
_sym = Symbol('x')
_symbols = Function('x')

#----------------------------------------------------------------------------#
#                                                                            #
#                           Fibonacci numbers                                #
#                                                                            #
#----------------------------------------------------------------------------#

class fibonacci(Function):
    r"""
    Fibonacci numbers / Fibonacci polynomials

    The Fibonacci numbers are the integer sequence defined by the
    initial terms F_0 = 0, F_1 = 1 and the two-term recurrence
    relation F_n = F_{n-1} + F_{n-2}.  This definition
    extended to arbitrary real and complex arguments using
    the formula

    .. math :: F_z = \frac{\phi^z - \cos(\pi z) \phi^{-z}}{\sqrt 5}

    The Fibonacci polynomials are defined by F_1(x) = 1,
    F_2(x) = x, and F_n(x) = x*F_{n-1}(x) + F_{n-2}(x) for n > 2.
    For all positive integers n, F_n(1) = F_n.

    * fibonacci(n) gives the nth Fibonacci number, F_n
    * fibonacci(n, x) gives the nth Fibonacci polynomial in x, F_n(x)

    Examples
    ========

    >>> from sympy import fibonacci, Symbol

    >>> [fibonacci(x) for x in range(11)]
    [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
    >>> fibonacci(5, Symbol('t'))
    t**4 + 3*t**2 + 1

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Fibonacci_number
    .. [2] http://mathworld.wolfram.com/FibonacciNumber.html

    See Also
    ========

    bell, bernoulli, catalan, euler, harmonic, lucas
    """

    @staticmethod
    def _fib(n):
        return _ifib(n)

    @staticmethod
    @recurrence_memo([None, S.One, _sym])
    def _fibpoly(n, prev):
        return (prev[-2] + _sym*prev[-1]).expand()

    @classmethod
    def eval(cls, n, sym=None):
        if n is S.Infinity:
            return S.Infinity

        if n.is_Integer:
            n = int(n)
            if n < 0:
                return S.NegativeOne**(n + 1) * fibonacci(-n)
            if sym is None:
                return Integer(cls._fib(n))
            else:
                if n < 1:
                    raise ValueError("Fibonacci polynomials are defined "
                       "only for positive integer indices.")
                return cls._fibpoly(n).subs(_sym, sym)

    def _eval_rewrite_as_sqrt(self, n):
        return 2**(-n)*sqrt(5)*((1 + sqrt(5))**n - (-sqrt(5) + 1)**n) / 5

    def _eval_rewrite_as_GoldenRatio(self,n):
        return (S.GoldenRatio**n - 1/(-S.GoldenRatio)**n)/(2*S.GoldenRatio-1)

class lucas(Function):
    """
    Lucas numbers

    Lucas numbers satisfy a recurrence relation similar to that of
    the Fibonacci sequence, in which each term is the sum of the
    preceding two. They are generated by choosing the initial
    values L_0 = 2 and L_1 = 1.

    * lucas(n) gives the nth Lucas number

    Examples
    ========

    >>> from sympy import lucas

    >>> [lucas(x) for x in range(11)]
    [2, 1, 3, 4, 7, 11, 18, 29, 47, 76, 123]

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Lucas_number
    .. [2] http://mathworld.wolfram.com/LucasNumber.html

    See Also
    ========

    bell, bernoulli, catalan, euler, fibonacci, harmonic
    """

    @classmethod
    def eval(cls, n):
        if n is S.Infinity:
            return S.Infinity

        if n.is_Integer:
            return fibonacci(n + 1) + fibonacci(n - 1)

    def _eval_rewrite_as_sqrt(self, n):
        return 2**(-n)*((1 + sqrt(5))**n + (-sqrt(5) + 1)**n)

#----------------------------------------------------------------------------#
#                                                                            #
#                           Bernoulli numbers                                #
#                                                                            #
#----------------------------------------------------------------------------#

class bernoulli(Function):
    r"""
    Bernoulli numbers / Bernoulli polynomials

    The Bernoulli numbers are a sequence of rational numbers
    defined by B_0 = 1 and the recursive relation (n > 0)::

                n
               ___
              \      / n + 1 \
          0 =  )     |       | * B .
              /___   \   k   /    k
              k = 0

    They are also commonly defined by their exponential generating
    function, which is x/(exp(x) - 1). For odd indices > 1, the
    Bernoulli numbers are zero.

    The Bernoulli polynomials satisfy the analogous formula::

                    n
                   ___
                  \      / n \         n-k
          B (x) =  )     |   | * B  * x   .
           n      /___   \ k /    k
                  k = 0

    Bernoulli numbers and Bernoulli polynomials are related as
    B_n(0) = B_n.

    We compute Bernoulli numbers using Ramanujan's formula::

                                   / n + 3 \
          B   =  (A(n) - S(n))  /  |       |
           n                       \   n   /

    where A(n) = (n+3)/3 when n = 0 or 2 (mod 6), A(n) = -(n+3)/6
    when n = 4 (mod 6), and::

                 [n/6]
                  ___
                 \      /  n + 3  \
          S(n) =  )     |         | * B
                 /___   \ n - 6*k /    n-6*k
                 k = 1

    This formula is similar to the sum given in the definition, but
    cuts 2/3 of the terms. For Bernoulli polynomials, we use the
    formula in the definition.

    * bernoulli(n) gives the nth Bernoulli number, B_n
    * bernoulli(n, x) gives the nth Bernoulli polynomial in x, B_n(x)

    Examples
    ========

    >>> from sympy import bernoulli

    >>> [bernoulli(n) for n in range(11)]
    [1, -1/2, 1/6, 0, -1/30, 0, 1/42, 0, -1/30, 0, 5/66]
    >>> bernoulli(1000001)
    0

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Bernoulli_number
    .. [2] http://en.wikipedia.org/wiki/Bernoulli_polynomial
    .. [3] http://mathworld.wolfram.com/BernoulliNumber.html
    .. [4] http://mathworld.wolfram.com/BernoulliPolynomial.html

    See Also
    ========

    bell, catalan, euler, fibonacci, harmonic, lucas
    """

    # Calculates B_n for positive even n
    @staticmethod
    def _calc_bernoulli(n):
        s = 0
        a = int(binomial(n + 3, n - 6))
        for j in range(1, n//6 + 1):
            s += a * bernoulli(n - 6*j)
            # Avoid computing each binomial coefficient from scratch
            a *= _product(n - 6 - 6*j + 1, n - 6*j)
            a //= _product(6*j + 4, 6*j + 9)
        if n % 6 == 4:
            s = -Rational(n + 3, 6) - s
        else:
            s = Rational(n + 3, 3) - s
        return s / binomial(n + 3, n)

    # We implement a specialized memoization scheme to handle each
    # case modulo 6 separately
    _cache = {0: S.One, 2: Rational(1, 6), 4: Rational(-1, 30)}
    _highest = {0: 0, 2: 2, 4: 4}

    @classmethod
    def eval(cls, n, sym=None):
        if n.is_Number:
            if n.is_Integer and n.is_nonnegative:
                if n is S.Zero:
                    return S.One
                elif n is S.One:
                    if sym is None:
                        return -S.Half
                    else:
                        return sym - S.Half
                # Bernoulli numbers
                elif sym is None:
                    if n.is_odd:
                        return S.Zero
                    n = int(n)
                    # Use mpmath for enormous Bernoulli numbers
                    if n > 500:
                        p, q = bernfrac(n)
                        return Rational(int(p), int(q))
                    case = n % 6
                    highest_cached = cls._highest[case]
                    if n <= highest_cached:
                        return cls._cache[n]
                    # To avoid excessive recursion when, say, bernoulli(1000) is
                    # requested, calculate and cache the entire sequence ... B_988,
                    # B_994, B_1000 in increasing order
                    for i in range(highest_cached + 6, n + 6, 6):
                        b = cls._calc_bernoulli(i)
                        cls._cache[i] = b
                        cls._highest[case] = i
                    return b
                # Bernoulli polynomials
                else:
                    n, result = int(n), []
                    for k in range(n + 1):
                        result.append(binomial(n, k)*cls(k)*sym**(n - k))
                    return Add(*result)
            else:
                raise ValueError("Bernoulli numbers are defined only"
                                 " for nonnegative integer indices.")

        if sym is None:
            if n.is_odd and (n - 1).is_positive:
                return S.Zero

#----------------------------------------------------------------------------#
#                                                                            #
#                             Bell numbers                                   #
#                                                                            #
#----------------------------------------------------------------------------#

class bell(Function):
    r"""
    Bell numbers / Bell polynomials

    The Bell numbers satisfy `B_0 = 1` and

    .. math:: B_n = \sum_{k=0}^{n-1} \binom{n-1}{k} B_k.

    They are also given by:

    .. math:: B_n = \frac{1}{e} \sum_{k=0}^{\infty} \frac{k^n}{k!}.

    The Bell polynomials are given by `B_0(x) = 1` and

    .. math:: B_n(x) = x \sum_{k=1}^{n-1} \binom{n-1}{k-1} B_{k-1}(x).

    The second kind of Bell polynomials (are sometimes called "partial" Bell
    polynomials or incomplete Bell polynomials) are defined as

    .. math:: B_{n,k}(x_1, x_2,\dotsc x_{n-k+1}) =
            \sum_{j_1+j_2+j_2+\dotsb=k \atop j_1+2j_2+3j_2+\dotsb=n}
                \frac{n!}{j_1!j_2!\dotsb j_{n-k+1}!}
                \left(\frac{x_1}{1!} \right)^{j_1}
                \left(\frac{x_2}{2!} \right)^{j_2} \dotsb
                \left(\frac{x_{n-k+1}}{(n-k+1)!} \right) ^{j_{n-k+1}}.

    * bell(n) gives the `n^{th}` Bell number, `B_n`.
    * bell(n, x) gives the `n^{th}` Bell polynomial, `B_n(x)`.
    * bell(n, k, (x1, x2, ...)) gives Bell polynomials of the second kind,
      `B_{n,k}(x_1, x_2, \dotsc, x_{n-k+1})`.

    Notes
    =====

    Not to be confused with Bernoulli numbers and Bernoulli polynomials,
    which use the same notation.

    Examples
    ========

    >>> from sympy import bell, Symbol, symbols

    >>> [bell(n) for n in range(11)]
    [1, 1, 2, 5, 15, 52, 203, 877, 4140, 21147, 115975]
    >>> bell(30)
    846749014511809332450147
    >>> bell(4, Symbol('t'))
    t**4 + 6*t**3 + 7*t**2 + t
    >>> bell(6, 2, symbols('x:6')[1:])
    6*x1*x5 + 15*x2*x4 + 10*x3**2

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Bell_number
    .. [2] http://mathworld.wolfram.com/BellNumber.html
    .. [3] http://mathworld.wolfram.com/BellPolynomial.html

    See Also
    ========

    bernoulli, catalan, euler, fibonacci, harmonic, lucas
    """

    @staticmethod
    @recurrence_memo([1, 1])
    def _bell(n, prev):
        s = 1
        a = 1
        for k in range(1, n):
            a = a * (n - k) // k
            s += a * prev[k]
        return s

    @staticmethod
    @recurrence_memo([S.One, _sym])
    def _bell_poly(n, prev):
        s = 1
        a = 1
        for k in range(2, n + 1):
            a = a * (n - k + 1) // (k - 1)
            s += a * prev[k - 1]
        return expand_mul(_sym * s)

    @staticmethod
    def _bell_incomplete_poly(n, k, symbols):
        r"""
        The second kind of Bell polynomials (incomplete Bell polynomials).

        Calculated by recurrence formula:

        .. math:: B_{n,k}(x_1, x_2, \dotsc, x_{n-k+1}) =
                \sum_{m=1}^{n-k+1}
                \x_m \binom{n-1}{m-1} B_{n-m,k-1}(x_1, x_2, \dotsc, x_{n-m-k})

        where
            B_{0,0} = 1;
            B_{n,0} = 0; for n>=1
            B_{0,k} = 0; for k>=1

        """
        if (n == 0) and (k == 0):
            return S.One
        elif (n == 0) or (k == 0):
            return S.Zero
        s = S.Zero
        a = S.One
        for m in range(1, n - k + 2):
            s += a * bell._bell_incomplete_poly(
                n - m, k - 1, symbols) * symbols[m - 1]
            a = a * (n - m) / m
        return expand_mul(s)

    @classmethod
    def eval(cls, n, k_sym=None, symbols=None):
        if n.is_Integer and n.is_nonnegative:
            if k_sym is None:
                return Integer(cls._bell(int(n)))
            elif symbols is None:
                return cls._bell_poly(int(n)).subs(_sym, k_sym)
            else:
                r = cls._bell_incomplete_poly(int(n), int(k_sym), symbols)
                return r

    def _eval_rewrite_as_Sum(self, n, k_sym=None, symbols=None):
        from sympy import Sum
        if (k_sym is not None) or (symbols is not None):
            return self

        # Dobinski's formula
        if not n.is_nonnegative:
            return self
        k = Dummy('k', integer=True, nonnegative=True)
        return 1 / E * Sum(k**n / factorial(k), (k, 0, S.Infinity))

#----------------------------------------------------------------------------#
#                                                                            #
#                           Harmonic numbers                                 #
#                                                                            #
#----------------------------------------------------------------------------#

class harmonic(Function):
    r"""
    Harmonic numbers

    The nth harmonic number is given by `\operatorname{H}_{n} =
    1 + \frac{1}{2} + \frac{1}{3} + \ldots + \frac{1}{n}`.

    More generally:

    .. math:: \operatorname{H}_{n,m} = \sum_{k=1}^{n} \frac{1}{k^m}

    As `n \rightarrow \infty`, `\operatorname{H}_{n,m} \rightarrow \zeta(m)`,
    the Riemann zeta function.

    * ``harmonic(n)`` gives the nth harmonic number, `\operatorname{H}_n`

    * ``harmonic(n, m)`` gives the nth generalized harmonic number
      of order `m`, `\operatorname{H}_{n,m}`, where
      ``harmonic(n) == harmonic(n, 1)``

    Examples
    ========

    >>> from sympy import harmonic, oo

    >>> [harmonic(n) for n in range(6)]
    [0, 1, 3/2, 11/6, 25/12, 137/60]
    >>> [harmonic(n, 2) for n in range(6)]
    [0, 1, 5/4, 49/36, 205/144, 5269/3600]
    >>> harmonic(oo, 2)
    pi**2/6

    >>> from sympy import Symbol, Sum
    >>> n = Symbol("n")

    >>> harmonic(n).rewrite(Sum)
    Sum(1/_k, (_k, 1, n))

    We can evaluate harmonic numbers for all integral and positive
    rational arguments:

    >>> from sympy import S, expand_func, simplify
    >>> harmonic(8)
    761/280
    >>> harmonic(11)
    83711/27720

    >>> H = harmonic(1/S(3))
    >>> H
    harmonic(1/3)
    >>> He = expand_func(H)
    >>> He
    -log(6) - sqrt(3)*pi/6 + 2*Sum(log(sin(_k*pi/3))*cos(2*_k*pi/3), (_k, 1, 1))
                           + 3*Sum(1/(3*_k + 1), (_k, 0, 0))
    >>> He.doit()
    -log(6) - sqrt(3)*pi/6 - log(sqrt(3)/2) + 3
    >>> H = harmonic(25/S(7))
    >>> He = simplify(expand_func(H).doit())
    >>> He
    log(sin(pi/7)**(-2*cos(pi/7))*sin(2*pi/7)**(2*cos(16*pi/7))*cos(pi/14)**(-2*sin(pi/14))/14)
    + pi*tan(pi/14)/2 + 30247/9900
    >>> He.n(40)
    1.983697455232980674869851942390639915940
    >>> harmonic(25/S(7)).n(40)
    1.983697455232980674869851942390639915940

    We can rewrite harmonic numbers in terms of polygamma functions:

    >>> from sympy import digamma, polygamma
    >>> m = Symbol("m")

    >>> harmonic(n).rewrite(digamma)
    polygamma(0, n + 1) + EulerGamma

    >>> harmonic(n).rewrite(polygamma)
    polygamma(0, n + 1) + EulerGamma

    >>> harmonic(n,3).rewrite(polygamma)
    polygamma(2, n + 1)/2 - polygamma(2, 1)/2

    >>> harmonic(n,m).rewrite(polygamma)
    (-1)**m*(polygamma(m - 1, 1) - polygamma(m - 1, n + 1))/factorial(m - 1)

    Integer offsets in the argument can be pulled out:

    >>> from sympy import expand_func

    >>> expand_func(harmonic(n+4))
    harmonic(n) + 1/(n + 4) + 1/(n + 3) + 1/(n + 2) + 1/(n + 1)

    >>> expand_func(harmonic(n-4))
    harmonic(n) - 1/(n - 1) - 1/(n - 2) - 1/(n - 3) - 1/n

    Some limits can be computed as well:

    >>> from sympy import limit, oo

    >>> limit(harmonic(n), n, oo)
    oo

    >>> limit(harmonic(n, 2), n, oo)
    pi**2/6

    >>> limit(harmonic(n, 3), n, oo)
    -polygamma(2, 1)/2

    However we can not compute the general relation yet:

    >>> limit(harmonic(n, m), n, oo)
    harmonic(oo, m)

    which equals ``zeta(m)`` for ``m > 1``.

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Harmonic_number
    .. [2] http://functions.wolfram.com/GammaBetaErf/HarmonicNumber/
    .. [3] http://functions.wolfram.com/GammaBetaErf/HarmonicNumber2/

    See Also
    ========

    bell, bernoulli, catalan, euler, fibonacci, lucas
    """

    # Generate one memoized Harmonic number-generating function for each
    # order and store it in a dictionary
    _functions = {}

    @classmethod
    def eval(cls, n, m=None):
        from sympy import zeta
        if m is S.One:
            return cls(n)
        if m is None:
            m = S.One

        if m.is_zero:
            return n

        if n is S.Infinity and m.is_Number:
            # TODO: Fix for symbolic values of m
            if m.is_negative:
                return S.NaN
            elif LessThan(m, S.One):
                return S.Infinity
            elif StrictGreaterThan(m, S.One):
                return zeta(m)
            else:
                return cls

        if n.is_Integer and n.is_nonnegative and m.is_Integer:
            if n == 0:
                return S.Zero
            if not m in cls._functions:
                @recurrence_memo([0])
                def f(n, prev):
                    return prev[-1] + S.One / n**m
                cls._functions[m] = f
            return cls._functions[m](int(n))

    def _eval_rewrite_as_polygamma(self, n, m=1):
        from sympy.functions.special.gamma_functions import polygamma
        return S.NegativeOne**m/factorial(m - 1) * (polygamma(m - 1, 1) - polygamma(m - 1, n + 1))

    def _eval_rewrite_as_digamma(self, n, m=1):
        from sympy.functions.special.gamma_functions import polygamma
        return self.rewrite(polygamma)

    def _eval_rewrite_as_trigamma(self, n, m=1):
        from sympy.functions.special.gamma_functions import polygamma
        return self.rewrite(polygamma)

    def _eval_rewrite_as_Sum(self, n, m=None):
        from sympy import Sum
        k = Dummy("k", integer=True)
        if m is None:
            m = S.One
        return Sum(k**(-m), (k, 1, n))

    def _eval_expand_func(self, **hints):
        from sympy import Sum
        n = self.args[0]
        m = self.args[1] if len(self.args) == 2 else 1

        if m == S.One:
            if n.is_Add:
                off = n.args[0]
                nnew = n - off
                if off.is_Integer and off.is_positive:
                    result = [S.One/(nnew + i) for i in range(off, 0, -1)] + [harmonic(nnew)]
                    return Add(*result)
                elif off.is_Integer and off.is_negative:
                    result = [-S.One/(nnew + i) for i in range(0, off, -1)] + [harmonic(nnew)]
                    return Add(*result)

            if n.is_Rational:
                # Expansions for harmonic numbers at general rational arguments (u + p/q)
                # Split n as u + p/q with p < q
                p, q = n.as_numer_denom()
                u = p // q
                p = p - u * q
                if u.is_nonnegative and p.is_positive and q.is_positive and p < q:
                    k = Dummy("k")
                    t1 = q * Sum(1 / (q * k + p), (k, 0, u))
                    t2 = 2 * Sum(cos((2 * pi * p * k) / S(q)) *
                                   log(sin((pi * k) / S(q))),
                                   (k, 1, floor((q - 1) / S(2))))
                    t3 = (pi / 2) * cot((pi * p) / q) + log(2 * q)
                    return t1 + t2 - t3

        return self

    def _eval_rewrite_as_tractable(self, n, m=1):
        from sympy import polygamma
        return self.rewrite(polygamma).rewrite("tractable", deep=True)

    def _eval_evalf(self, prec):
        from sympy import polygamma
        if all(i.is_number for i in self.args):
            return self.rewrite(polygamma)._eval_evalf(prec)

#----------------------------------------------------------------------------#
#                                                                            #
#                           Euler numbers                                    #
#                                                                            #
#----------------------------------------------------------------------------#

class euler(Function):
    r"""
    Euler numbers

    The euler numbers are given by::

                  2*n+1   k
                   ___   ___            j          2*n+1
                  \     \     / k \ (-1)  * (k-2*j)
          E   = I  )     )    |   | --------------------
           2n     /___  /___  \ j /      k    k
                  k = 1 j = 0           2  * I  * k

          E     = 0
           2n+1

    * euler(n) gives the n-th Euler number, E_n

    Examples
    ========

    >>> from sympy import Symbol
    >>> from sympy.functions import euler
    >>> [euler(n) for n in range(10)]
    [1, 0, -1, 0, 5, 0, -61, 0, 1385, 0]
    >>> n = Symbol("n")
    >>> euler(n+2*n)
    euler(3*n)

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Euler_numbers
    .. [2] http://mathworld.wolfram.com/EulerNumber.html
    .. [3] http://en.wikipedia.org/wiki/Alternating_permutation
    .. [4] http://mathworld.wolfram.com/AlternatingPermutation.html

    See Also
    ========

    bell, bernoulli, catalan, fibonacci, harmonic, lucas
    """

    @classmethod
    def eval(cls, m):
        if m.is_odd:
            return S.Zero
        if m.is_Integer and m.is_nonnegative:
            from mpmath import mp
            m = m._to_mpmath(mp.prec)
            res = mp.eulernum(m, exact=True)
            return Integer(res)

    def _eval_rewrite_as_Sum(self, arg):
        from sympy import Sum
        if arg.is_even:
            k = Dummy("k", integer=True)
            j = Dummy("j", integer=True)
            n = self.args[0] / 2
            Em = (S.ImaginaryUnit * Sum(Sum(binomial(k, j) * ((-1)**j * (k - 2*j)**(2*n + 1)) /
                  (2**k*S.ImaginaryUnit**k * k), (j, 0, k)), (k, 1, 2*n + 1)))

            return Em

    def _eval_evalf(self, prec):
        m = self.args[0]

        if m.is_Integer and m.is_nonnegative:
            from mpmath import mp
            from sympy import Expr
            m = m._to_mpmath(prec)
            with workprec(prec):
                res = mp.eulernum(m)
            return Expr._from_mpmath(res, prec)

#----------------------------------------------------------------------------#
#                                                                            #
#                           Catalan numbers                                  #
#                                                                            #
#----------------------------------------------------------------------------#

class catalan(Function):
    r"""
    Catalan numbers

    The n-th catalan number is given by::

                 1   / 2*n \
          C  = ----- |     |
           n   n + 1 \  n  /

    * catalan(n) gives the n-th Catalan number, C_n

    Examples
    ========

    >>> from sympy import (Symbol, binomial, gamma, hyper, polygamma,
    ...             catalan, diff, combsimp, Rational, I)

    >>> [ catalan(i) for i in range(1,10) ]
    [1, 2, 5, 14, 42, 132, 429, 1430, 4862]

    >>> n = Symbol("n", integer=True)

    >>> catalan(n)
    catalan(n)

    Catalan numbers can be transformed into several other, identical
    expressions involving other mathematical functions

    >>> catalan(n).rewrite(binomial)
    binomial(2*n, n)/(n + 1)

    >>> catalan(n).rewrite(gamma)
    4**n*gamma(n + 1/2)/(sqrt(pi)*gamma(n + 2))

    >>> catalan(n).rewrite(hyper)
    hyper((-n + 1, -n), (2,), 1)

    For some non-integer values of n we can get closed form
    expressions by rewriting in terms of gamma functions:

    >>> catalan(Rational(1,2)).rewrite(gamma)
    8/(3*pi)

    We can differentiate the Catalan numbers C(n) interpreted as a
    continuous real funtion in n:

    >>> diff(catalan(n), n)
    (polygamma(0, n + 1/2) - polygamma(0, n + 2) + log(4))*catalan(n)

    As a more advanced example consider the following ratio
    between consecutive numbers:

    >>> combsimp((catalan(n + 1)/catalan(n)).rewrite(binomial))
    2*(2*n + 1)/(n + 2)

    The Catalan numbers can be generalized to complex numbers:

    >>> catalan(I).rewrite(gamma)
    4**I*gamma(1/2 + I)/(sqrt(pi)*gamma(2 + I))

    and evaluated with arbitrary precision:

    >>> catalan(I).evalf(20)
    0.39764993382373624267 - 0.020884341620842555705*I

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Catalan_number
    .. [2] http://mathworld.wolfram.com/CatalanNumber.html
    .. [3] http://functions.wolfram.com/GammaBetaErf/CatalanNumber/
    .. [4] http://geometer.org/mathcircles/catalan.pdf

    See Also
    ========

    bell, bernoulli, euler, fibonacci, harmonic, lucas
    sympy.functions.combinatorial.factorials.binomial
    """

    @classmethod
    def eval(cls, n):
        from sympy import gamma
        if (n.is_Integer and n.is_nonnegative) or \
           (n.is_noninteger and n.is_negative):
            return 4**n*gamma(n + S.Half)/(gamma(S.Half)*gamma(n + 2))

        if (n.is_integer and n.is_negative):
            if (n + 1).is_negative:
                return S.Zero
            if (n + 1).is_zero:
                return -S.Half

    def fdiff(self, argindex=1):
        from sympy import polygamma, log
        n = self.args[0]
        return catalan(n)*(polygamma(0, n + Rational(1, 2)) - polygamma(0, n + 2) + log(4))

    def _eval_rewrite_as_binomial(self, n):
        return binomial(2*n, n)/(n + 1)

    def _eval_rewrite_as_factorial(self, n):
        return factorial(2*n) / (factorial(n+1) * factorial(n))

    def _eval_rewrite_as_gamma(self, n):
        from sympy import gamma
        # The gamma function allows to generalize Catalan numbers to complex n
        return 4**n*gamma(n + S.Half)/(gamma(S.Half)*gamma(n + 2))

    def _eval_rewrite_as_hyper(self, n):
        from sympy import hyper
        return hyper([1 - n, -n], [2], 1)

    def _eval_rewrite_as_Product(self, n):
        from sympy import Product
        if not (n.is_integer and n.is_nonnegative):
            return self
        k = Dummy('k', integer=True, positive=True)
        return Product((n + k) / k, (k, 2, n))

    def _eval_evalf(self, prec):
        from sympy import gamma
        if self.args[0].is_number:
            return self.rewrite(gamma)._eval_evalf(prec)

#----------------------------------------------------------------------------#
#                                                                            #
#                           Genocchi numbers                                 #
#                                                                            #
#----------------------------------------------------------------------------#

class genocchi(Function):
    r"""
    Genocchi numbers

    The Genocchi numbers are a sequence of integers G_n that satisfy the
    relation::

                           oo
                         ____
                         \   `
                 2*t      \         n
                ------ =   \   G_n*t
                 t         /   ------
                e  + 1    /      n!
                         /___,
                         n = 1

    Examples
    ========

    >>> from sympy import Symbol
    >>> from sympy.functions import genocchi
    >>> [genocchi(n) for n in range(1, 9)]
    [1, -1, 0, 1, 0, -3, 0, 17]
    >>> n = Symbol('n', integer=True, positive=True)
    >>> genocchi(2 * n + 1)
    0

    References
    ==========

    .. [1] https://en.wikipedia.org/wiki/Genocchi_number
    .. [2] http://mathworld.wolfram.com/GenocchiNumber.html

    See Also
    ========

    bell, bernoulli, catalan, euler, fibonacci, harmonic, lucas
    """

    @classmethod
    def eval(cls, n):
        if n.is_Number:
            if (not n.is_Integer) or n.is_nonpositive:
                raise ValueError("Genocchi numbers are defined only for " +
                                 "positive integers")
            return 2 * (1 - S(2) ** n) * bernoulli(n)

        if n.is_odd and (n - 1).is_positive:
            return S.Zero

        if (n - 1).is_zero:
            return S.One

    def _eval_rewrite_as_bernoulli(self, n):
        if n.is_integer and n.is_nonnegative:
            return (1 - S(2) ** n) * bernoulli(n) * 2

    def _eval_is_integer(self):
        if self.args[0].is_integer and self.args[0].is_positive:
            return True

    def _eval_is_negative(self):
        n = self.args[0]
        if n.is_integer and n.is_positive:
            if n.is_odd:
                return False
            return (n / 2).is_odd

    def _eval_is_positive(self):
        n = self.args[0]
        if n.is_integer and n.is_positive:
            if n.is_odd:
                return fuzzy_not((n - 1).is_positive)
            return (n / 2).is_even

    def _eval_is_even(self):
        n = self.args[0]
        if n.is_integer and n.is_positive:
            if n.is_even:
                return False
            return (n - 1).is_positive

    def _eval_is_odd(self):
        n = self.args[0]
        if n.is_integer and n.is_positive:
            if n.is_even:
                return True
            return fuzzy_not((n - 1).is_positive)

    def _eval_is_prime(self):
        n = self.args[0]
        # only G_6 = -3 and G_8 = 17 are prime,
        # but SymPy does not consider negatives as prime
        # so only n=8 is tested
        return (n - 8).is_zero

#######################################################################
###
### Functions for enumerating partitions, permutations and combinations
###
#######################################################################

class _MultisetHistogram(tuple):
    pass

_N = -1
_ITEMS = -2
_M = slice(None, _ITEMS)

def _multiset_histogram(n):
    """Return tuple used in permutation and combination counting. Input
    is a dictionary giving items with counts as values or a sequence of
    items (which need not be sorted).

    The data is stored in a class deriving from tuple so it is easily
    recognized and so it can be converted easily to a list.
    """
    if type(n) is dict:  # item: count
        if not all(isinstance(v, int) and v >= 0 for v in n.values()):
            raise ValueError
        tot = sum(n.values())
        items = sum(1 for k in n if n[k] > 0)
        return _MultisetHistogram([n[k] for k in n if n[k] > 0] + [items, tot])
    else:
        n = list(n)
        s = set(n)
        if len(s) == len(n):
            n = [1]*len(n)
            n.extend([len(n), len(n)])
            return _MultisetHistogram(n)
        m = dict(zip(s, range(len(s))))
        d = dict(zip(range(len(s)), [0]*len(s)))
        for i in n:
            d[m[i]] += 1
        return _multiset_histogram(d)

def nP(n, k=None, replacement=False):
    """Return the number of permutations of ``n`` items taken ``k`` at a time.

    Possible values for ``n``::
        integer - set of length ``n``
        sequence - converted to a multiset internally
        multiset - {element: multiplicity}

    If ``k`` is None then the total of all permutations of length 0
    through the number of items represented by ``n`` will be returned.

    If ``replacement`` is True then a given item can appear more than once
    in the ``k`` items. (For example, for 'ab' permutations of 2 would
    include 'aa', 'ab', 'ba' and 'bb'.) The multiplicity of elements in
    ``n`` is ignored when ``replacement`` is True but the total number
    of elements is considered since no element can appear more times than
    the number of elements in ``n``.

    Examples
    ========

    >>> from sympy.functions.combinatorial.numbers import nP
    >>> from sympy.utilities.iterables import multiset_permutations, multiset
    >>> nP(3, 2)
    6
    >>> nP('abc', 2) == nP(multiset('abc'), 2) == 6
    True
    >>> nP('aab', 2)
    3
    >>> nP([1, 2, 2], 2)
    3
    >>> [nP(3, i) for i in range(4)]
    [1, 3, 6, 6]
    >>> nP(3) == sum(_)
    True

    When ``replacement`` is True, each item can have multiplicity
    equal to the length represented by ``n``:

    >>> nP('aabc', replacement=True)
    121
    >>> [len(list(multiset_permutations('aaaabbbbcccc', i))) for i in range(5)]
    [1, 3, 9, 27, 81]
    >>> sum(_)
    121

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Permutation

    See Also
    ========
    sympy.utilities.iterables.multiset_permutations

    """
    try:
        n = as_int(n)
    except ValueError:
        return Integer(_nP(_multiset_histogram(n), k, replacement))
    return Integer(_nP(n, k, replacement))

@cacheit
def _nP(n, k=None, replacement=False):
    from sympy.functions.combinatorial.factorials import factorial
    from sympy.core.mul import prod

    if k == 0:
        return 1
    if isinstance(n, SYMPY_INTS):  # n different items
        # assert n >= 0
        if k is None:
            return sum(_nP(n, i, replacement) for i in range(n + 1))
        elif replacement:
            return n**k
        elif k > n:
            return 0
        elif k == n:
            return factorial(k)
        elif k == 1:
            return n
        else:
            # assert k >= 0
            return _product(n - k + 1, n)
    elif isinstance(n, _MultisetHistogram):
        if k is None:
            return sum(_nP(n, i, replacement) for i in range(n[_N] + 1))
        elif replacement:
            return n[_ITEMS]**k
        elif k == n[_N]:
            return factorial(k)/prod([factorial(i) for i in n[_M] if i > 1])
        elif k > n[_N]:
            return 0
        elif k == 1:
            return n[_ITEMS]
        else:
            # assert k >= 0
            tot = 0
            n = list(n)
            for i in range(len(n[_M])):
                if not n[i]:
                    continue
                n[_N] -= 1
                if n[i] == 1:
                    n[i] = 0
                    n[_ITEMS] -= 1
                    tot += _nP(_MultisetHistogram(n), k - 1)
                    n[_ITEMS] += 1
                    n[i] = 1
                else:
                    n[i] -= 1
                    tot += _nP(_MultisetHistogram(n), k - 1)
                    n[i] += 1
                n[_N] += 1
            return tot

@cacheit
def _AOP_product(n):
    """for n = (m1, m2, .., mk) return the coefficients of the polynomial,
    prod(sum(x**i for i in range(nj + 1)) for nj in n); i.e. the coefficients
    of the product of AOPs (all-one polynomials) or order given in n.  The
    resulting coefficient corresponding to x**r is the number of r-length
    combinations of sum(n) elements with multiplicities given in n.
    The coefficients are given as a default dictionary (so if a query is made
    for a key that is not present, 0 will be returned).

    Examples
    ========

    >>> from sympy.functions.combinatorial.numbers import _AOP_product
    >>> from sympy.abc import x
    >>> n = (2, 2, 3)  # e.g. aabbccc
    >>> prod = ((x**2 + x + 1)*(x**2 + x + 1)*(x**3 + x**2 + x + 1)).expand()
    >>> c = _AOP_product(n); dict(c)
    {0: 1, 1: 3, 2: 6, 3: 8, 4: 8, 5: 6, 6: 3, 7: 1}
    >>> [c[i] for i in range(8)] == [prod.coeff(x, i) for i in range(8)]
    True

    The generating poly used here is the same as that listed in
    http://tinyurl.com/cep849r, but in a refactored form.

    """
    from collections import defaultdict

    n = list(n)
    ord = sum(n)
    need = (ord + 2)//2
    rv = [1]*(n.pop() + 1)
    rv.extend([0]*(need - len(rv)))
    rv = rv[:need]
    while n:
        ni = n.pop()
        N = ni + 1
        was = rv[:]
        for i in range(1, min(N, len(rv))):
            rv[i] += rv[i - 1]
        for i in range(N, need):
            rv[i] += rv[i - 1] - was[i - N]
    rev = list(reversed(rv))
    if ord % 2:
        rv = rv + rev
    else:
        rv[-1:] = rev
    d = defaultdict(int)
    for i in range(len(rv)):
        d[i] = rv[i]
    return d

def nC(n, k=None, replacement=False):
    """Return the number of combinations of ``n`` items taken ``k`` at a time.

    Possible values for ``n``::
        integer - set of length ``n``
        sequence - converted to a multiset internally
        multiset - {element: multiplicity}

    If ``k`` is None then the total of all combinations of length 0
    through the number of items represented in ``n`` will be returned.

    If ``replacement`` is True then a given item can appear more than once
    in the ``k`` items. (For example, for 'ab' sets of 2 would include 'aa',
    'ab', and 'bb'.) The multiplicity of elements in ``n`` is ignored when
    ``replacement`` is True but the total number of elements is considered
    since no element can appear more times than the number of elements in
    ``n``.

    Examples
    ========

    >>> from sympy.functions.combinatorial.numbers import nC
    >>> from sympy.utilities.iterables import multiset_combinations
    >>> nC(3, 2)
    3
    >>> nC('abc', 2)
    3
    >>> nC('aab', 2)
    2

    When ``replacement`` is True, each item can have multiplicity
    equal to the length represented by ``n``:

    >>> nC('aabc', replacement=True)
    35
    >>> [len(list(multiset_combinations('aaaabbbbcccc', i))) for i in range(5)]
    [1, 3, 6, 10, 15]
    >>> sum(_)
    35

    If there are ``k`` items with multiplicities ``m_1, m_2, ..., m_k``
    then the total of all combinations of length 0 hrough ``k`` is the
    product, ``(m_1 + 1)*(m_2 + 1)*...*(m_k + 1)``. When the multiplicity
    of each item is 1 (i.e., k unique items) then there are 2**k
    combinations. For example, if there are 4 unique items, the total number
    of combinations is 16:

    >>> sum(nC(4, i) for i in range(5))
    16

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Combination
    .. [2] http://tinyurl.com/cep849r

    See Also
    ========
    sympy.utilities.iterables.multiset_combinations
    """
    from sympy.functions.combinatorial.factorials import binomial
    from sympy.core.mul import prod

    if isinstance(n, SYMPY_INTS):
        if k is None:
            if not replacement:
                return 2**n
            return sum(nC(n, i, replacement) for i in range(n + 1))
        if k < 0:
            raise ValueError("k cannot be negative")
        if replacement:
            return binomial(n + k - 1, k)
        return binomial(n, k)
    if isinstance(n, _MultisetHistogram):
        N = n[_N]
        if k is None:
            if not replacement:
                return prod(m + 1 for m in n[_M])
            return sum(nC(n, i, replacement) for i in range(N + 1))
        elif replacement:
            return nC(n[_ITEMS], k, replacement)
        # assert k >= 0
        elif k in (1, N - 1):
            return n[_ITEMS]
        elif k in (0, N):
            return 1
        return _AOP_product(tuple(n[_M]))[k]
    else:
        return nC(_multiset_histogram(n), k, replacement)

@cacheit
def _stirling1(n, k):
    if n == k == 0:
        return S.One
    if 0 in (n, k):
        return S.Zero
    n1 = n - 1

    # some special values
    if n == k:
        return S.One
    elif k == 1:
        return factorial(n1)
    elif k == n1:
        return binomial(n, 2)
    elif k == n - 2:
        return (3*n - 1)*binomial(n, 3)/4
    elif k == n - 3:
        return binomial(n, 2)*binomial(n, 4)

    # general recurrence
    return n1*_stirling1(n1, k) + _stirling1(n1, k - 1)

@cacheit
def _stirling2(n, k):
    if n == k == 0:
        return S.One
    if 0 in (n, k):
        return S.Zero
    n1 = n - 1

    # some special values
    if k == n1:
        return binomial(n, 2)
    elif k == 2:
        return 2**n1 - 1

    # general recurrence
    return k*_stirling2(n1, k) + _stirling2(n1, k - 1)

def stirling(n, k, d=None, kind=2, signed=False):
    """Return Stirling number S(n, k) of the first or second (default) kind.

    The sum of all Stirling numbers of the second kind for k = 1
    through n is bell(n). The recurrence relationship for these numbers
    is::

    {0}       {n}   {0}      {n + 1}     {n}   {  n  }
    { } = 1;  { } = { } = 0; {     } = j*{ } + {     }
    {0}       {0}   {k}      {  k  }     {k}   {k - 1}

    where ``j`` is::
        ``n`` for Stirling numbers of the first kind
        ``-n`` for signed Stirling numbers of the first kind
        ``k`` for Stirling numbers of the second kind

    The first kind of Stirling number counts the number of permutations of
    ``n`` distinct items that have ``k`` cycles; the second kind counts the
    ways in which ``n`` distinct items can be partitioned into ``k`` parts.
    If ``d`` is given, the "reduced Stirling number of the second kind" is
    returned: ``S^{d}(n, k) = S(n - d + 1, k - d + 1)`` with ``n >= k >= d``.
    (This counts the ways to partition ``n`` consecutive integers into
    ``k`` groups with no pairwise difference less than ``d``. See example
    below.)

    To obtain the signed Stirling numbers of the first kind, use keyword
    ``signed=True``. Using this keyword automatically sets ``kind`` to 1.

    Examples
    ========

    >>> from sympy.functions.combinatorial.numbers import stirling, bell
    >>> from sympy.combinatorics import Permutation
    >>> from sympy.utilities.iterables import multiset_partitions, permutations

    First kind (unsigned by default):

    >>> [stirling(6, i, kind=1) for i in range(7)]
    [0, 120, 274, 225, 85, 15, 1]
    >>> perms = list(permutations(range(4)))
    >>> [sum(Permutation(p).cycles == i for p in perms) for i in range(5)]
    [0, 6, 11, 6, 1]
    >>> [stirling(4, i, kind=1) for i in range(5)]
    [0, 6, 11, 6, 1]

    First kind (signed):

    >>> [stirling(4, i, signed=True) for i in range(5)]
    [0, -6, 11, -6, 1]

    Second kind:

    >>> [stirling(10, i) for i in range(12)]
    [0, 1, 511, 9330, 34105, 42525, 22827, 5880, 750, 45, 1, 0]
    >>> sum(_) == bell(10)
    True
    >>> len(list(multiset_partitions(range(4), 2))) == stirling(4, 2)
    True

    Reduced second kind:

    >>> from sympy import subsets, oo
    >>> def delta(p):
    ...    if len(p) == 1:
    ...        return oo
    ...    return min(abs(i[0] - i[1]) for i in subsets(p, 2))
    >>> parts = multiset_partitions(range(5), 3)
    >>> d = 2
    >>> sum(1 for p in parts if all(delta(i) >= d for i in p))
    7
    >>> stirling(5, 3, 2)
    7

    References
    ==========

    .. [1] http://en.wikipedia.org/wiki/Stirling_numbers_of_the_first_kind
    .. [2] http://en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind

    See Also
    ========
    sympy.utilities.iterables.multiset_partitions

    """
    # TODO: make this a class like bell()

    n = as_int(n)
    k = as_int(k)
    if n < 0:
        raise ValueError('n must be nonnegative')
    if k > n:
        return S.Zero
    if d:
        # assert k >= d
        # kind is ignored -- only kind=2 is supported
        return _stirling2(n - d + 1, k - d + 1)
    elif signed:
        # kind is ignored -- only kind=1 is supported
        return (-1)**(n - k)*_stirling1(n, k)

    if kind == 1:
        return _stirling1(n, k)
    elif kind == 2:
        return _stirling2(n, k)
    else:
        raise ValueError('kind must be 1 or 2, not %s' % k)

@cacheit
def _nT(n, k):
    """Return the partitions of ``n`` items into ``k`` parts. This
    is used by ``nT`` for the case when ``n`` is an integer."""
    if k == 0:
        return 1 if k == n else 0
    return sum(_nT(n - k, j) for j in range(min(k, n - k) + 1))

def nT(n, k=None):
    """Return the number of ``k``-sized partitions of ``n`` items.

    Possible values for ``n``::
        integer - ``n`` identical items
        sequence - converted to a multiset internally
        multiset - {element: multiplicity}

    Note: the convention for ``nT`` is different than that of ``nC`` and
    ``nP`` in that
    here an integer indicates ``n`` *identical* items instead of a set of
    length ``n``; this is in keeping with the ``partitions`` function which
    treats its integer-``n`` input like a list of ``n`` 1s. One can use
    ``range(n)`` for ``n`` to indicate ``n`` distinct items.

    If ``k`` is None then the total number of ways to partition the elements
    represented in ``n`` will be returned.

    Examples
    ========

    >>> from sympy.functions.combinatorial.numbers import nT

    Partitions of the given multiset:

    >>> [nT('aabbc', i) for i in range(1, 7)]
    [1, 8, 11, 5, 1, 0]
    >>> nT('aabbc') == sum(_)
    True

    >>> [nT("mississippi", i) for i in range(1, 12)]
    [1, 74, 609, 1521, 1768, 1224, 579, 197, 50, 9, 1]

    Partitions when all items are identical:

    >>> [nT(5, i) for i in range(1, 6)]
    [1, 2, 2, 1, 1]
    >>> nT('1'*5) == sum(_)
    True

    When all items are different:

    >>> [nT(range(5), i) for i in range(1, 6)]
    [1, 15, 25, 10, 1]
    >>> nT(range(5)) == sum(_)
    True

    References
    ==========

    .. [1] http://undergraduate.csse.uwa.edu.au/units/CITS7209/partition.pdf

    See Also
    ========
    sympy.utilities.iterables.partitions
    sympy.utilities.iterables.multiset_partitions

    """
    from sympy.utilities.enumerative import MultisetPartitionTraverser

    if isinstance(n, SYMPY_INTS):
        # assert n >= 0
        # all the same
        if k is None:
            return sum(_nT(n, k) for k in range(1, n + 1))
        return _nT(n, k)
    if not isinstance(n, _MultisetHistogram):
        try:
            # if n contains hashable items there is some
            # quick handling that can be done
            u = len(set(n))
            if u == 1:
                return nT(len(n), k)
            elif u == len(n):
                n = range(u)
            raise TypeError
        except TypeError:
            n = _multiset_histogram(n)
    N = n[_N]
    if k is None and N == 1:
        return 1
    if k in (1, N):
        return 1
    if k == 2 or N == 2 and k is None:
        m, r = divmod(N, 2)
        rv = sum(nC(n, i) for i in range(1, m + 1))
        if not r:
            rv -= nC(n, m)//2
        if k is None:
            rv += 1  # for k == 1
        return rv
    if N == n[_ITEMS]:
        # all distinct
        if k is None:
            return bell(N)
        return stirling(N, k)
    m = MultisetPartitionTraverser()
    if k is None:
        return m.count_partitions(n[_M])
    # MultisetPartitionTraverser does not have a range-limited count
    # method, so need to enumerate and count
    tot = 0
    for discard in m.enum_range(n[_M], k-1, k):
        tot += 1
    return tot


# -*- coding: utf-8 -*-
'''
Stage: "シャーロックの家"
'''
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '../../..'))
sys.path.append('storybuilder')
from storybuilder.builder.world import World

# NOTE
#   物語が展開する中心。事務所
#   シャーロックが借りて住んでいる古い一軒家（二階建て）で、B221という看板が近くにある
#   幽霊が出るといういわくつきの物件で安かった。原因は地場発生の魔石があっただけ
#   ：２F
#   [物置][空き部屋]
#   [寝室１][寝室２]
#   ：１F
#   [キッチン][バス]
#   [ダイニング][トイレ]
#   [研究室]
#   [リビング][書斎]

# alias
HOME = "SherlockHouse"
LIVING = "SherlockHouseLiving"
KITCHEN = "SherlockHouseKitchen"
BATHROOM = "SherlockHouseBathroom"
DINING = "SherlockHouseDining"
LABO = "SherlockHouseLabo"
READING = "SherlockHouseReadingRoom"
BEDROOM = "SherlockHouseBedroom"
STORAGE = "SherlockHouseStorage"

## scenes
def about_sherlock(w: World):
    return w.scene("$sherlockという男について",
            w.plot_note("知っていたならなぜ忠告してくれなかったんだ、と$wilsonが$sherlockに文句をいい、なんとか家に入れてもらえる"),
            w.plot_note("中に入るとそこら中に本や資料がちらばっていた"),
            w.plot_note("$wilsonはそこで依頼をしようと思ったが、"),
            w.plot_note("そこに何も記載のない手紙が投げ込まれた"),
            )

def read_prince_letter(w: World):
    return w.scene("皇太子の書簡の中身",
            w.plot_note("手紙には独特の紙が使われていて、それが王室のものだと$sherlockは分かった"),
            w.plot_note("中は皇太子からの手紙で、$sherlockに頼みごとが書かれていた"),
            w.plot_note("皇太子は女遊びがひどくてその界隈では有名だが、今回ついに腰を落ち着けて結婚することになった"),
            w.plot_note("相手は近隣の公国の王女で、政治的な意味合いも大きい"),
            w.plot_note("その結婚に際して過去の女性関係をすべて綺麗にした"),
            w.plot_note("ただある一人の女性にプレゼントしてしまった大切なナイフを返してもらいたいが、相手の女性が応じてくれない"),
            w.plot_note("揉め事をおこしたくないので、穏便にすませたいから、$sherlockに彼女を説得して、ナイフを返してもらってくれないか、という依頼"),
            w.plot_note("$sherlockはその依頼内容について、書かれていない部分の推測を述べる"),
            w.plot_note("ナイフと書いているが、実際は王室に伝わる宝剣で、それが王の証の一つで、結婚の際には儀式内で使われる"),
            w.plot_note("酒の勢いで大切な宝剣をあげてしまったのだろうと"),
            w.plot_note("そんなものを取り戻す義理はないが、恩があるので仕方なく依頼を受けると言った"),
            w.plot_note("$sherlockは$wilsonにその女性の家まで送ってほしいと頼んだ"),
            )

def important_than_sword(w: World):
    return w.scene("宝剣より大事なこと",
            w.plot_note("$sherlockは宝剣よりも殺人事件についての調査をしたいと、$wilsonを家に置いて出ていってしまう"),
            w.plot_note("$wilsonは$sherlockの家に戻り、そこで彼を待つことにする"),
            w.plot_note("やってきた若い刑事は$sherlockがいないことに落胆しつつも、状況を教えてくれる"),
            w.plot_note("発見された遺体は一月ほど前に行方不明になった女性だった"),
            w.plot_note("$ailyとは何の関係もなく、そこの接点も見つけられないと嘆く"),
            w.plot_note("殺害方法も不明で、凶器すら見つけられないと"),
            w.plot_note("そこに役所の男から$ailyという女性が住民登録をしたという形跡は見つけられなかったと連絡がきた"),
            )

def housemate_mary(w: World):
    return w.scene("同居人$mary",
            w.plot_note("同居するようになった$maryはやたらと$sherlockにまとわりつく"),
            w.plot_note("$sherlockは大好きな読書もできず、困っていた"),
            )

def mary_has_worry(w: World):
    return w.scene("$maryの悩み",
            w.plot_note("$maryは彼の迷惑になりたくなくて、$wilsonに相談する"),
            w.plot_note("女手が不足しているから自分が役立つところをアピールしてみたら、と助言を受ける"),
            w.plot_note("$maryは掃除や買い物を買って出る"),
            )

def about_missings(w: World):
    return w.scene("失踪者について",
            w.plot_note("やっと外に出てくれてほっとした$sherlockは$wilsonに事件について相談する"),
            w.plot_note("最近謎の失踪者が増えていた"),
            w.plot_note("失踪事件として新聞や雑誌も特集を組んでいる"),
            w.plot_note("$wilsonはその調査を$sherlockに依頼していたが、未だに何も情報がなかった"),
            )

def strange_armor_knight(w: World):
    return w.scene("奇妙な鎧騎士",
            w.plot_note("そこに$maryが見知らぬ人を連れて戻ってくる"),
            w.plot_note("道端で困っていたから拾ったけれど言葉がしゃべれないのだと$maryは説明した"),
            )

def strange_work(w: World):
    return w.scene("奇妙な仕事",
            w.plot_note("その鎧騎士は$sherlockに$limeと名乗った（筆談で）"),
            w.plot_note("彼女は今ある老夫婦の家に居候しているが、彼らの知人の質屋の護衛のアルバイトをしていた"),
            w.plot_note("守衛仲間の$binsと交代しながら閉店時刻まで警備をしている"),
            w.plot_note("その$binsから別のバイトを紹介され、今は途中にそちらもやっている"),
            w.plot_note("その別のバイトが相談したいことだった"),
            w.plot_note("最初に$binsからチラシを見せてもらったときには「赤い鎧の者だけがバイト資格がある」と書かれていた"),
            w.plot_note("仕事内容はじっと座ってある本の写しを作る作業を三時間行うだけで、週給で結構な金額がもらえた"),
            w.plot_note("実際に面接に行ってみると確かに赤い鎧を来た人間が集まっていたが、$limeみたいに見事に全身赤という者はいなかった"),
            w.plot_note("主催者である赤鎧クラブは彼女を合格とし、その翌日から守衛を抜け出して三時間、そのアルバイトをしているらしい"),
            w.plot_note("オーナー夫婦には申し訳なく感じているが、そのお金でプレゼントしたいと思っているのだと説明する"),
            w.plot_note("その話をきいて$sherlockは彼女に今すぐそのアルバイトを辞めるようにとだけ言った"),
            )

def reason_for_lime_work(w: World):
    return w.scene("$sherlockの忠告",
            w.plot_note("家に帰った$maryはどうしてあんな風に言ったのか$sherlockに問いただす"),
            w.plot_note("$sherlockはそんなにうまい話は存在しないし、自分が知る限り「赤鎧クラブ」なんてものは存在しないと断言する"),
            w.plot_note("$maryは実際に持ち帰ったチラシを見せながら、彼女を拾ってくれたオーナーさんや同僚の$binsの優しさを力説する"),
            w.plot_note("しかし後日$sherlockの言っていたように問題が発生する"),
            w.plot_note("その近所にあった改装中の銀行が強盗に襲われた"),
            w.plot_note("警備員が気づいて連絡したが、表からも裏からも誰も入ってはおらず、謎の強盗と話題になっていた"),
            w.plot_note("しかし現地を調べたところ、抜け穴が掘ってあり、大量のダイヤと金塊が盗まれたあとだった"),
            w.plot_note("しかもその抜け穴は質屋に繋がっていたのだ"),
            w.plot_note("その質屋のオーナー夫婦も逮捕され、$limeも容疑者の一人として逮捕された"),
            )

def help_lime_please(w: World):
    return w.scene("$limeを助けて",
            w.plot_note("$maryが$limeを助けてやってほしいと$sherlockに言う"),
            w.plot_note("$sherlockは自分の忠告を聞かなかったからだと言うが、それでも話だけは聞くと言う"),
            w.plot_note("質屋につながっていた抜け穴の中で、重要参考人だった$ignesが遺体で発見された"),
            w.plot_note("その容疑者として$limeが逮捕され、オーナー夫婦も事情聴取を受けている最中らしい"),
            w.plot_note("強盗の件についても調査中で、全部彼女に押し付けられるかもしれないと言い出す"),
            w.plot_note("$sherlockはその質屋に案内してもらう"),
            )

def limes_talk_of_strange_case(w: World):
    return w.scene("奇妙な事件についての$limeの話",
            w.plot_note("$maryが$limeを拾い、再び家へと連れてくる"),
            w.plot_note("$sherlockは銀行から盗まれたものがダイヤだけじゃないと睨むが、教えてもらえなかった"),
            w.plot_note("家に戻ってくると$sherlockはそこに$limeがいることに頭を抱える"),
            w.plot_note("$limeがしゃべれないのは呪いの鎧のせいだと言う"),
            w.plot_note("その呪いをといてもらおうと、知人の神官を読んでいた"),
            w.plot_note("呪いを解いたが$limeはしゃべれないままだった"),
            w.plot_note("その$limeは筆談で自分が王室の人間であると告白する"),
            )

def lime_was_royal_family(w: World):
    return w.scene("$limeは王家の人間",
            w.plot_note("$limeは自分が誘拐された訳ではなく、普通に家出をしたのだと告白する"),
            w.plot_note("王室はそんな品の悪い発表をできないから失踪事件にして公表したのだと言った"),
            w.plot_note("もともと妾の子で、周囲から浮いていて、王室にも自分の居場所がなく帰りたくないと泣く"),
            )

def newcommer_lime(w: World):
    return w.scene("新しい同居人$lime",
            w.plot_note("$maryは$sherlockに$limeを一緒に住まわせてほしいとお願いする"),
            w.plot_note("$sherlockは金銭的な問題さえ解決できればと提案する"),
            w.plot_note("$wilsonは金のことなら大丈夫だと、なぜか大金を手にして言う"),
            w.plot_note("$wilsonは$sherlockの秘蔵コレクションを売り払っていた"),
            w.plot_note("こうして新しい住人$limeをここに加えることになった"),
            )

def cooker_lime(w: World):
    return w.scene("料理人$lime",
            w.plot_note("$limeは料理担当になっていて、そのガチョウをもらってさばいてくれる"),
            )

def marys_market_talk(w: World):
    return w.scene("$maryの市場の話",
            w.plot_note("$maryは市場で仕入れた面白い話を$sherlockに話す"),
            w.plot_note("今市場ではガチョウからダイヤが出てくると話題になっていた"),
            w.plot_note("$limeがやってきて、何か出たという"),
            w.plot_note("ガチョウの中から出てきたのは血がついたナイフだった"),
            )

def knife_in_the_goose(w: World):
    return w.scene("ガチョウの中の凶器",
            w.plot_note("$sherlockはそれがなにかの事件の凶器だと分かり、すぐに警察に連絡を取る"),
            w.plot_note("$restradeがやってきて、それは現在彼が追っている事件の重要な証拠品だと言われた"),
            )

def restrade_talk_about_goose_knife(w: World):
    return w.scene("$restradeのガチョウの凶器事件の話",
            w.plot_note("$restradeからその事件についての概要を聞く"),
            w.plot_note("事件はある一軒家で起こった"),
            w.plot_note("引退した学者が謎の死を遂げた"),
            w.plot_note("刺殺だったのだが凶器が発見されず、犯人も特定されないまま現在に至る"),
            w.plot_note("そのナイフを警部に渡して調べてもらう"),
            w.plot_note("その間に興味をもった$sherlockは一人でその現場を調べに出ていってしまう"),
            w.plot_note("後日、そのナイフからずっと失踪中の$jackの指紋が検出された"),
            )

def backhome_mary_with_jack_wanted(w: World):
    return w.scene("$jack容疑者の話を持って返ってきた$mary",
            w.plot_note("戻ってきた$sherlockは$maryからそのことを聞き、"),
            w.plot_note("$sherlockは現場を見てきたことを$maryたちに話す"),
            )

def talk_about_goose_case(w: World):
    return w.scene("ガチョウ凶器事件についての調査",
            w.plot_note("現場は住宅街から少し離れた郊外の一軒家で、男は民間の研究所をやめたあとも個人的に何かを研究していた"),
            w.plot_note("歴史学と民俗学に造形が深く、$sherlockもその所蔵していた資料に関心をしたくらい"),
            w.plot_note("彼が書き残しているものの一つに古代の技法がいくつか紹介されていた"),
            w.plot_note("刺された場所は彼の家だが、凶器は消えている"),
            w.plot_note("ただし$jackとの関係性は全く見えず、彼女ならそんな手段を使わないと$sherlockは考えた"),
            w.plot_note("$sherlockは誰かが$jackを表舞台に引っ張り出したい、その罠だと考える"),
            )

def jacks_letter(w: World):
    return w.scene("$jackからの手紙",
            w.plot_note("と、差出人不明の手紙に$jackからのメッセージがあった"),
            w.plot_note("助けてほしいと"),
            )

def sherlocks_message_for_jack(w: World):
    return w.scene("$sherlockのメッセージ",
            w.plot_note("そこに$sherlockからの伝言を$ignesが持ってくる"),
            w.plot_note("数日留守にすることと、$jackに会いに行ってくると書かれていた"),
            )

def mysterious_case(w :World):
    return w.scene("怪奇事件",
            w.plot_note("$sherlockは怪奇事件の特集記事を読みながら「こんなものは実在しない」と言う"),
            w.plot_note("そもそも奇妙な現象、霊的なもの、不思議なものは人間が理解することを放棄していると説明する"),
            w.plot_note("小さい頃、学校内で七不思議というものがあったが、それを全て解明したらみんなから怒られたと"),
            )

def legend_of_dark_dog(w: World):
    return w.scene("魔獣伝説",
            w.plot_note("そこに$wilsonがこんな話がある、と、ある孤島に伝わる魔獣伝説を話した"),
            w.plot_note("そこはこの三年の間に六名もの犠牲者が出ているという"),
            w.plot_note("最初は飼い犬や家畜が殺されているだけだった"),
            w.plot_note("しかし最初に人の犠牲者が出た"),
            w.plot_note("それはどう見ても人の手によるものではなく、何か獣による被害だった"),
            w.plot_note("最初の事件から次の事件まではかなり時間が開いたが、直近はこの三ヶ月の間に二件も殺人事件が起こっている"),
            w.plot_note("$sherlockはそれだけ続くなら必ず人の手が関わっていると断言する"),
            )

def invitation_from_dark_island(w: World):
    return w.scene("孤島からの招待状",
            w.plot_note("そこに招待状が届く"),
            w.plot_note("$wilsonはそれを開封し、噂をしていれば、とその伝説の孤島に暮らす城主からの招待状だと言った"),
            )

def commision_of_murder_case(w: World):
    return w.scene("殺人事件の解決依頼",
            w.plot_note("$sherlockは新聞を読んでいた"),
            w.plot_note("そこに殺人事件の調査依頼が持ち込まれる"),
            w.plot_note("最初は$maryも驚いていたが今では慣れたもので、依頼人を案内して、飲み物を出しながら依頼内容を話すよう促す"),
            w.plot_note("$maryは秘書気取りだった"),
            w.plot_note("だが$sherlockは依頼人が出した名前に驚く"),
            w.plot_note("それは$morianoの大学の後輩だったからだ"),
            w.plot_note("犯罪学の研究をしている人間が殺された"),
            w.plot_note("大学の研究室内での密室殺人。その手口が全く不明だが自殺ではないと警察は断定しているという"),
            w.plot_note("さっそくその調査に向かう$sherlock"),
            )

def moriano_is_here(w: World):
    return w.scene("$moriano見参",
            w.plot_note("$sherlockが家に戻ってくるとそこには老人の姿があった", "$morianoだ"),
            w.plot_note("$morianoは「はじめまして」と挨拶をし、それから今$sherlockたちがどういう経路で戻ってきたかを言い当てる"),
            w.plot_note("$morianoは$sherlockに自分に関するすべてのことから手を引くようにと警告する"),
            w.plot_note("$sherlockは$morianoがここに来ることも推測して既に逮捕する準備を整えているとブラフを張るが、彼には通用しなかった"),
            w.plot_note("警察は別のところで起こった事件に駆けつけている"),
            w.plot_note("$morianoは言う。すべての人間は自分の意志ではなく、環境要因によって動かされると。つまり誰でもが犯罪者になりうると"),
            w.plot_note("$morianoは$maryに問いかける。彼女は$sherlockを好きだろうと"),
            w.plot_note("$morianoは$limeに本心では王室に帰りたいだろうと"),
            w.plot_note("$wilsonについての言及はとくないが、ここの人間には言えない本音を隠しているだろうと"),
            w.plot_note("$morianoは逃げないからいつでも自分の屋敷に来るがいいと言い残して、去っていく"),
            )

def marys_strange(w: World):
    return w.scene("奇妙な$mary",
            w.plot_note("$morianoがきてから$maryの様子がおかしい"),
            w.plot_note("$sherlockは$morianoを何とか見つけ出そうと躍起になっている"),
            w.plot_note("$maryは$limeに相談することもできず、市場の$nowlisに愚痴る。自分だけが違う気がすると"),
            )

def where_is_mary(w: World):
    return w.scene("$maryはどこへ？",
            w.plot_note("すぐに$sherlockたちは$morianoの邸宅に向かう"),
            )

def alive_moriano(w: World):
    return w.scene("$morianoは生きている",
            w.plot_note("しかし後日、$morianoのメッセージが新聞に掲載される"),
            w.plot_note("$maryは無事で丁重に監禁していると。場所は$sherlockなら推理できると書かれて、ヒントが残されていた"),
            )

def morianos_whereabouts(w: World):
    return w.scene("$morianoの居場所",
            w.plot_note("$sherlockはヒントから$maryの居場所は$morianoと関係ない場所にいると推測する"),
            w.plot_note("$mary救出隊として少年探偵団の協力を仰ぐ"),
            w.plot_note("その間に$sherlockはその新聞記事からたどり、$morianoがどこからメッセージを出しているのかを調べる"),
            )

def vanished_sherlock(w: World):
    return w.scene("消えた$sherlock",
            w.plot_note("$maryたちが戻ると、そこには$sherlockの姿がなかった"),
            )

def sherlocks_information(w: World):
    return w.scene("$sherlockに関する情報",
            w.plot_note("$maryが目覚めるとそこに$sherlockの姿がいなかった"),
            w.plot_note("戻ってきた$wilsonは$sherlockの手がかりを追ったが見失ったと言う"),
            w.plot_note("$maryは$sherlockが戻ってくると信じて待っていたが、連絡も戻ってくることもなかった"),
            )

def no_sherlock_life(w: World):
    return w.scene("$sherlock不在の生活",
            "同・寝室",
            w.plot_note("一月が経ち、$maryたちは$sherlockのいない生活に馴染み始めていた"),
            )

def serching_sherlock(w: World):
    return w.scene("$sherlockの捜索隊",
            w.plot_note("町では$morianoも$sherlockも消えたというのに犯罪は起こっていた"),
            w.plot_note("$wilsonは手を尽くして$sherlockを探す"),
            w.plot_note("$limeが王室のツテを使い、何とか情報を集めると言い出す"),
            w.plot_note("少年探偵団も手を尽くした"),
            w.plot_note("今まで世話になった人たちも$sherlockのことを探してくれた"),
            w.plot_note("それでも情報すら見つからない"),
            )

def arrived_his_message(w: World):
    return w.scene("$sherlockからのメッセージ",
            w.plot_note("雨の酷いある日、一通の郵便が届く"),
            w.plot_note("届いた手紙には宛名がなかったが$sherlock特有の署名が入っていた"),
            )

def sadness_report(w: World):
    return w.scene("悲しいお知らせ",
            w.plot_note("手紙の冒頭にはこう書かれていた"),
            w.plot_note("この手紙が届いたならば自分はすでにこの世界にいないだろうと$sherlockは書いていた"),
            w.plot_note("手紙は$morianoの隠れ家に向かう直前に書いて出したと書かれている"),
            w.plot_note("$morianoは用意周到で、約束通り一人で待っていたりはしない"),
            w.plot_note("$sherlockは陽動をして、手下たちを遠ざけて、可能ならば$morianoをおびき出す"),
            w.plot_note("どうにか一対一で話せる場所を作る、と書いてある"),
            w.plot_note("$morianoがどこまで語るか、告白するかわからないが、彼がやってきた悪事について書き残しておく"),
            w.plot_note("$morianoのことは以前教えたが、大学を出たあとの彼については今回独自調査を行うまで不明な部分が多かった"),
            w.plot_note("$morianoは$cultXと呼ばれる宗教団体との接触から犯罪者人生が始まる"),
            w.plot_note("彼はその教義である人間の本性である「悪」を反映させようとしていた"),
            w.plot_note("今までに解決した事件の裏側にはこの教団か、その教団の人間、関係者が細い糸で繋がっていた"),
            w.plot_note("その大本である$morianoを何としても打ち取ると宣言されていた"),
            w.plot_note("$maryたちは$sherlockがどうなったのか気になり、手紙を出した場所に向かおうとする"),
            w.plot_note("だが$wilsonによりそれは止められる"),
            w.plot_note("兄の$mikelがやってきて「$sherlockがなくなった」と告げた"),
            )

## in Empty House
def believed_his_alive(w: World):
    mary, lime, wil = w.get("mary"), w.get("lime"), w.get("wilson")
    return w.scene("$sherlockの生存を信じて",
            w.change_camera("mary"),
            w.change_stage(KITCHEN),
            w.change_time("morning"),
            w.plot_note("$maryたちは$sherlockが生きていると思って捜索を続けていた"),
            w.plot_note("しかし何の情報もなく、ただ時間だけが過ぎていく"),
            w.plot_note("家を失い、$wilsonの住まいに居候していた$maryたち"),
            w.plot_note("$wilsonは忙しそうに外に出ていることが増えた"),
            w.plot_note("$maryは$sherlockの手紙にヒントはないかと考えるが、何も見つからない"),
            mary.be("皿洗いをしている$S"),
            mary.think("もう一月も$sherlockは失踪を続けている"),
            mary.think("完全に死んだものと思われていたが、$Sたちは捜索を続けていた"),
            mary.do("棚には$sherlockのコップが残っている"),
            mary.think("それを目にして涙が滲む"),
            mary.think("でも$sherlockが送ってきたメッセージにはわざわざ自分が死んだと思ってくれと書かれていた"),
            mary.think("何故そんなことを書いたのか、$Sは気になっていた"),
            mary.talk("あっ"),
            mary.do("$wilsonの湯呑が割れてしまう"),
            mary.talk("$wilsonのだし、いいか"),
            )

def news_of_sherlock_alive(w: World):
    mary, lime, wil = w.get("mary"), w.get("lime"), w.get("wilson")
    return w.scene("$sherlock生存情報",
            w.change_camera("mary"),
            w.change_stage(LIVING),
            w.change_time("noon"),
            w.plot_note("だが$limeはそこに$sherlockが生きているという証拠を見つけた"),
            w.plot_note("そこに$wilsonが戻ってくる"),
            w.plot_note("$wilsonは「$sherlockに似た人間を見かけた」という情報を聞いたと話した"),
            mary.come("買い物を終えて帰ってきた$S"),
            lime.be("$Sは家の片付けをしていた"),
            wil.come("そこに$Sが興奮した様子で戻ってくる"),
            mary.talk("何かあったん？"),
            wil.talk("聞いてくれ", "いた", "$sherlockが、いたんだ"),
            mary.do("驚きで声が出ない$S"),
            lime.do("掃除の手が止まる$S"),
            wil.talk("$meもまだ聞いたばかりの話で、本当かどうかの確認すらできていないんだが、それでもこれまで何の情報もなかったところにこれは大きいよ"),
            wil.talk("$EastEndの空き家に夜な夜な明かりが灯る家があるそうなんだ",
                "どうやらそこに$sherlockによく似た人間が入っていくのを見たって、ホームレスの目撃情報があった"),
            mary.talk("でもどうしてそれが$sherlockなん？　別人の可能性はないん？"),
            wil.talk("それが以前$sherlockが世話をしたホームレスで、彼のことをよく覚えていたんだよ",
                "遠目にもあの特徴的な寝癖頭とそこに被ったハンチング、チェック柄のコートは$sherlockに間違いないって"),
            mary.think("その話に興奮する$S"),
            mary.talk("場所は？"),
            mary.do("荷物を置くと、$wilsonに詰め寄った"),
            )

def consideration_of_sherlock(w: World):
    return w.scene("容疑者$sherlockについての考察",
            # NOTE: omit?
            w.plot_note("一旦家に戻り、犯人にされてしまった$sherlockについて考える"),
            w.plot_note("$wilsonは$sherlockが$moriano一味に騙されたというのだが"),
            w.plot_note("もう一度あの空き家を訪れる"),
            )

def help_from_sherlock(w: World):
    lime = w.get("lime")
    wil = w.get("wilson")
    ignes = w.get("ignes")
    return w.scene("$sherlockからの救援情報",
            w.change_camera("lime"),
            w.change_stage(LIVING),
            lime.be("一人で$sherlockの家に戻っている$S"),
            lime.do("消えた$maryを探してくるとでかけた$wilson"),
            lime.do("$Sはポストに入っていた宛名のない封書を見つける"),
            lime.do("そこには$maryが$morianoの手の者に捕まり、監禁されていると書かれていた"),
            ignes.come("$Sがやってきて"),
            ignes.talk("$mary嬢ちゃんは？"),
            ignes.do("事情を聞く$S"),
            ignes.talk("すぐ手配して、場所を突き止める", "$limeさんは警察に行って事情を説明してきて"),
            )

def injured_wilson(w: World):
    return w.scene("負傷した$wilson",
            # NOTE: omit／方針変更
            w.plot_note("家に戻ると$wilsonがいて、ひどい怪我を負っていたが、無事に逃げ出したと言う"),
            w.plot_note("$maryは自分たちを助けた男が$sherlockの生存を言っていたと伝える"),
            w.plot_note("$wilsonはそのホームレスのことを教えてくれと頼む"),
            w.plot_note("$maryたちにここで休むようにいい、$wilsonは$sherlockを探しに出ていった"),
            w.plot_note("そこに$wilsonが指名手配されたと$restradeがやってくる"),
            )

def burned_shal_home(w: World):
    shal = w.get("sherlock")
    wil, lime = w.get("wilson"), w.get("lime")
    ignes, pat = w.get("ignes"), w.get("patson")
    lisa = w.get("lisa")
    return w.scene("家が燃えて",
            w.change_camera("sherlock"),
            w.change_stage(HOME),
            shal.come("$Sたちは$wilsonの車で火事で全焼してしまった住居前にやってくる"),
            wil.come(),
            lime.come(),
            shal.do("みんな呆然としてその光景を見ている"),
            shal.do("消防士たちが$magicポンプで水をかけている。少しずつ火勢は衰え、もう消火が近い"),
            shal.do("近所の人も出て、野次馬が集まっている"),
            lisa.come("大家の$Sがやってきて、びっくりして呆然"),
            lisa.talk("な、何なんですか、これは！"),
            shal.talk("ああ、$ln_lisaさん、どうもご無沙汰しています"),
            lisa.talk("ねえ$sherlockさん、これは一体どういうことなのかしら"),
            shal.talk("火事みたいですね。おそらく放火でしょう。迷惑なことです"),
            lisa.talk("燃えたのは誰の家なの？"),
            shal.talk("$meが借りていたあなたの家です"),
            lisa.talk("ええ、そうよね。そうだと思ったわ"),
            lisa.do("見る間に表情が変わっていく夫人"),
            lisa.talk("この弁償、誰が支払ってくれるのかしら"),
            shal.talk("契約上は$meに過失があった場合は$meですが、放火の責任を取れと言われても困りますから、おそらくオーナー夫人の方になるかと"),
            lisa.talk("何ですって！"),
            shal.talk("用事があったのを思い出したので失礼"),
            shal.do("$wilsonの車に乗り込む$S"),
            wil.do("仕方ない、といった感じで車に乗り込む$S"),
            )


from .bedloader import ROI


from __future__ import annotations

import ast
import functools
import sys
from typing import Iterable

from tokenize_rt import NON_CODING_TOKENS
from tokenize_rt import Offset
from tokenize_rt import Token

from pyupgrade._ast_helpers import ast_to_offset
from pyupgrade._ast_helpers import is_name_attr
from pyupgrade._data import register
from pyupgrade._data import State
from pyupgrade._data import TokenFunc
from pyupgrade._token_helpers import CLOSING
from pyupgrade._token_helpers import find_closing_bracket
from pyupgrade._token_helpers import find_token
from pyupgrade._token_helpers import OPENING

def _fix_optional(i: int, tokens: list[Token]) -> None:
    j = find_token(tokens, i, '[')
    k = find_closing_bracket(tokens, j)
    if tokens[j].line == tokens[k].line:
        tokens[k] = Token('CODE', ' | None')
        del tokens[i:j + 1]
    else:
        tokens[j] = tokens[j]._replace(src='(')
        tokens[k] = tokens[k]._replace(src=')')
        tokens[i:j] = [Token('CODE', 'None | ')]

def _fix_union(
        i: int,
        tokens: list[Token],
        *,
        arg_count: int,
) -> None:
    depth = 1
    parens_done = []
    open_parens = []
    commas = []
    coding_depth = None

    j = find_token(tokens, i, '[')
    k = j + 1
    while depth:
        # it's possible our first coding token is a close paren
        # so make sure this is separate from the if chain below
        if (
                tokens[k].name not in NON_CODING_TOKENS and
                tokens[k].src != '(' and
                coding_depth is None
        ):
            if tokens[k].src == ')':  # the coding token was an empty tuple
                coding_depth = depth - 1
            else:
                coding_depth = depth

        if tokens[k].src in OPENING:
            if tokens[k].src == '(':
                open_parens.append((depth, k))

            depth += 1
        elif tokens[k].src in CLOSING:
            if tokens[k].src == ')':
                paren_depth, open_paren = open_parens.pop()
                parens_done.append((paren_depth, (open_paren, k)))

            depth -= 1
        elif tokens[k].src == ',':
            commas.append((depth, k))

        k += 1
    k -= 1

    assert coding_depth is not None
    assert not open_parens, open_parens
    comma_depth = min((depth for depth, _ in commas), default=sys.maxsize)
    min_depth = min(comma_depth, coding_depth)

    to_delete = [
        paren
        for depth, positions in parens_done
        if depth < min_depth
        for paren in positions
    ]

    if comma_depth <= coding_depth:
        comma_positions = [k for depth, k in commas if depth == comma_depth]
        if len(comma_positions) == arg_count:
            to_delete.append(comma_positions.pop())
    else:
        comma_positions = []

    to_delete.sort()

    if tokens[j].line == tokens[k].line:
        del tokens[k]
        for comma in comma_positions:
            tokens[comma] = Token('CODE', ' |')
        for paren in reversed(to_delete):
            del tokens[paren]
        del tokens[i:j + 1]
    else:
        tokens[j] = tokens[j]._replace(src='(')
        tokens[k] = tokens[k]._replace(src=')')

        for comma in comma_positions:
            tokens[comma] = Token('CODE', ' |')
        for paren in reversed(to_delete):
            del tokens[paren]
        del tokens[i:j]

def _supported_version(state: State) -> bool:
    return (
        state.in_annotation and (
            state.settings.min_version >= (3, 10) or (
                not state.settings.keep_runtime_typing and
                'annotations' in state.from_imports['__future__']
            )
        )
    )

def _any_arg_is_str(node_slice: ast.expr) -> bool:
    return (
        isinstance(node_slice, ast.Str) or (
            isinstance(node_slice, ast.Tuple) and
            any(isinstance(elt, ast.Str) for elt in node_slice.elts)
        )
    )

@register(ast.Subscript)
def visit_Subscript(
        state: State,
        node: ast.Subscript,
        parent: ast.AST,
) -> Iterable[tuple[Offset, TokenFunc]]:
    if not _supported_version(state):
        return

    # prevent rewriting forward annotations
    if (
            (sys.version_info >= (3, 9) and _any_arg_is_str(node.slice)) or
            (
                sys.version_info < (3, 9) and
                isinstance(node.slice, ast.Index) and
                _any_arg_is_str(node.slice.value)
            )
    ):
        return

    if is_name_attr(
            node.value,
            state.from_imports,
            ('typing',),
            ('Optional',),
    ):
        yield ast_to_offset(node), _fix_optional
    elif is_name_attr(node.value, state.from_imports, ('typing',), ('Union',)):
        if sys.version_info >= (3, 9):  # pragma: >=3.9 cover
            node_slice = node.slice
        elif isinstance(node.slice, ast.Index):  # pragma: <3.9 cover
            node_slice: ast.AST = node.slice.value
        else:  # pragma: <3.9 cover
            node_slice = node.slice  # unexpected slice type

        if isinstance(node_slice, ast.Slice):  # not a valid annotation
            return

        if isinstance(node_slice, ast.Tuple):
            if node_slice.elts:
                arg_count = len(node_slice.elts)
            else:
                return  # empty Union
        else:
            arg_count = 1

        func = functools.partial(_fix_union, arg_count=arg_count)
        yield ast_to_offset(node), func


import argparse
import logging
import time

import sys
from twilio.rest import Client

import settings
import RPi.GPIO as GPIO

twilio = Client(settings.TWILIO_PUBLIC_KEY, settings.TWILIO_SECRET_KEY)
log = logging.getLogger(__name__)

class SaltLevelMonitor(object):
    def __init__(self, force_report=False, unit=settings.METRIC, threshold=0,
                 tank_depth=settings.DEFAULT_TANK_DEPTH):
        self.force_report = force_report
        self.unit = unit if unit in settings.VALID_UNITS else settings.METRIC
        self.notation = 'inches' if unit == settings.IMPERIAL else 'centimeters'
        self.threshold = float(threshold)
        self.tank_depth = float(tank_depth)
        self.distance = None
        self.remaining_salt = None

    def check_salt_level(self):
        self.distance = self.get_average_distance()
        self._convert_units()
        self.remaining_salt = self.tank_depth - self.distance
        message = self._get_report_message()
        log.info('Salt level is: {0:.2f} {1}'.format(self.remaining_salt, self.notation))
        if self.remaining_salt < self.threshold or self.force_report:
            log.info(message['body'])
            self.report_salt_level(message)

    def get_average_distance(self):
        """ used to get an average read since the sensor isn't 100% accurate """
        reads = [self.get_distance() for _ in range(settings.READS_PER_CHECK)]
        return sum(reads) / settings.READS_PER_CHECK

    @staticmethod
    def get_distance():
        """ returns distance in centimeters """
        # set Trigger to HIGH
        GPIO.output(settings.GPIO_TRIGGER, True)

        # set Trigger after 0.01ms to LOW
        time.sleep(0.00001)
        GPIO.output(settings.GPIO_TRIGGER, False)

        start_time = time.time()
        stop_time = time.time()

        # save StartTime
        while GPIO.input(settings.GPIO_ECHO) == 0:
            start_time = time.time()

        # save time of arrival
        while GPIO.input(settings.GPIO_ECHO) == 1:
            stop_time = time.time()

        # time difference between start and arrival
        time_elapsed = stop_time - start_time
        return (time_elapsed * settings.SPEED_OF_SOUND) / 2

    def _convert_units(self):
        """
        convert distance to inches if IMPERIAL or convert tank_depth and threshold to centimeters
        """
        if self.unit == settings.IMPERIAL:
            self.distance = self.distance / settings.CM_TO_INCHES
        else:
            self.tank_depth = self.tank_depth * settings.CM_TO_INCHES
            self.threshold = self.threshold * settings.CM_TO_INCHES

    def _get_report_message(self):
        message = settings.MESSAGE_TEMPLATE.copy()
        message['body'] = settings.SALT_LEVEL_ALERT_MESSAGE.format(
            self.remaining_salt, self.notation)
        if self.force_report:
            message['body'] = '{} (forced report)'.format(message['body'])
        return message

    @staticmethod
    def report_salt_level(message):
        twilio.messages.create(**message)

    def __enter__(self):
        GPIO.setmode(GPIO.BCM)

        # set GPIO direction (IN / OUT)
        GPIO.setup(settings.GPIO_TRIGGER, GPIO.OUT)
        GPIO.setup(settings.GPIO_ECHO, GPIO.IN)
        return self

    def __exit__(self, *args):
        GPIO.cleanup()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Salty Dog')
    parser.add_argument('-u',
                        '--unit',
                        action='store',
                        dest='unit',
                        default='metric',
                        help='Unit of measure used in reporting')
    parser.add_argument('-t',
                        '--threshold',
                        action='store',
                        dest='threshold',
                        help='Threshold for reporting in inches or cm (must match --unit)')
    parser.add_argument('-d',
                        '--tank-depth',
                        action='store',
                        dest='tank_depth',
                        help='Total depth of your salt tank in inches or cm (must match --unit)')
    parser.add_argument('-f',
                        '--force-report',
                        action='store_true',
                        dest='force_report',
                        default=False,
                        help='Force Salty Dog to send SMS regardless of salt level measured')
    args = parser.parse_args(sys.argv[1:])
    parsed_kwargs = {
        'force_report': args.force_report,
        'unit': args.unit,
        'threshold': args.threshold,
        'tank_depth': args.tank_depth,
    }
    with SaltLevelMonitor(**parsed_kwargs) as monitor:
        monitor.check_salt_level()