import unittest

import lib.encodings
import lib.ling

from . import tools

from .tools import (
    assert_equal,
    assert_false,
    assert_in,
    assert_is,
    assert_is_instance,
    assert_is_none,
    assert_not_equal,
    assert_not_in,
    assert_raises,
    assert_true,
    collect_yielded,
)

L = lib.ling
T = lib.ling.Language
E = lib.encodings

class test_fix_codes:

    def t(self, l1, l2):
        lang = L.parse_language(l1)
        assert_equal(str(lang), l1)
        if l1 == l2:
            assert_is_none(lang.fix_codes())
        else:
            assert_is(lang.fix_codes(), True)
        assert_equal(str(lang), l2)

    def test_2_to_2(self):
        self.t('grc', 'grc')
        self.t('grc_GR', 'grc_GR')

    def test_1_to_1(self):
        self.t('el', 'el')
        self.t('el_GR', 'el_GR')

    def test_2t_to_1(self):
        self.t('ell', 'el')
        self.t('ell_GR', 'el_GR')

    def test_2b_to_1(self):
        self.t('gre', 'el')
        self.t('gre_GR', 'el_GR')

    def test_ll_not_found(self):
        with assert_raises(L.FixingLanguageCodesFailed):
            self.t('ry', '')

    def test_cc_not_found(self):
        with assert_raises(L.FixingLanguageCodesFailed):
            self.t('el_RY', '')

def test_language_repr():
    # Language.__repr__() is never used by i18nspector itself,
    # but it's useful for debugging test failures.
    lng = T('el')
    assert_equal(repr(lng), '<Language el>')

class test_language_equality:

    # ==, !=, is_almost_equal()

    def test_eq(self):
        l1 = T('el', 'GR')
        l2 = T('el', 'GR')
        assert_equal(l1, l2)
        assert_equal(l2, l1)

    def test_ne(self):
        l1 = T('el')
        l2 = T('el', 'GR')
        assert_not_equal(l1, l2)
        assert_not_equal(l2, l1)

    def test_ne_other_type(self):
        l1 = T('el')
        assert_not_equal(l1, 42)
        assert_not_equal(42, l1)

    def test_almost_equal(self):
        l1 = T('el')
        l2 = T('el', 'GR')
        assert_true(l1.is_almost_equal(l2))
        assert_true(l2.is_almost_equal(l1))

    def test_not_almost_equal(self):
        l1 = T('el', 'GR')
        l2 = T('grc', 'GR')
        assert_false(l1.is_almost_equal(l2))
        assert_false(l2.is_almost_equal(l1))

    def test_not_almost_equal_other_type(self):
        l1 = T('el')
        with assert_raises(TypeError):
            l1.is_almost_equal(42)

class test_remove_encoding:

    def t(self, l1, l2):
        lang = L.parse_language(l1)
        assert_equal(str(lang), l1)
        if l1 == l2:
            assert_is_none(lang.remove_encoding())
        else:
            assert_is(lang.remove_encoding(), True)
        assert_equal(str(lang), l2)

    def test_without_encoding(self):
        self.t('el', 'el')

    def test_with_encoding(self):
        self.t('el.UTF-8', 'el')

class test_remove_nonlinguistic_modifier:

    def t(self, l1, l2):
        lang = L.parse_language(l1)
        assert_equal(str(lang), l1)
        if l1 == l2:
            assert_is_none(lang.remove_nonlinguistic_modifier())
        else:
            assert_is(lang.remove_nonlinguistic_modifier(), True)
        assert_equal(str(lang), l2)

    def test_quot(self):
        self.t('en@quot', 'en@quot')
        self.t('en@boldquot', 'en@boldquot')

    def test_latin(self):
        self.t('sr@latin', 'sr@latin')

    def test_euro(self):
        self.t('de_AT@euro', 'de_AT')

class test_lookup_territory_code:

    def test_found(self):
        cc = L.lookup_territory_code('GR')
        assert_equal(cc, 'GR')

    def test_not_found(self):
        cc = L.lookup_territory_code('RG')
        assert_is_none(cc)

class test_get_language_for_name:

    def t(self, name, expected):
        lang = L.get_language_for_name(name)
        assert_is_instance(lang, T)
        assert_equal(str(lang), expected)

    def test_found(self):
        self.t('Greek', 'el')

    def test_found_multi(self):
        self.t('Old Church Slavonic', 'cu')

    def test_found_as_ascii(self):
        self.t('Norwegian Bokmål', 'nb')

    def test_found_semicolon(self):
        self.t('Chichewa; Nyanja', 'ny')

    def test_found_comma(self):
        self.t('Ndebele, South', 'nr')

    def test_found_comma_as_semicolon(self):
        self.t('Pashto, Pushto', 'ps')

    def test_lone_comma(self):
        with assert_raises(LookupError):
            self.t(',', None)

    def test_not_found(self):
        with assert_raises(LookupError):
            self.t('Nadsat', None)

class test_parse_language:

    def test_ll(self):
        lang = L.parse_language('el')
        assert_equal(lang.language_code, 'el')
        assert_is_none(lang.territory_code)
        assert_is_none(lang.encoding)
        assert_is_none(lang.modifier)

    def test_lll(self):
        lang = L.parse_language('ell')
        assert_equal(lang.language_code, 'ell')
        assert_is_none(lang.territory_code)
        assert_is_none(lang.encoding)
        assert_is_none(lang.modifier)

    def test_ll_cc(self):
        lang = L.parse_language('el_GR')
        assert_equal(lang.language_code, 'el')
        assert_equal(lang.territory_code, 'GR')
        assert_is_none(lang.encoding)
        assert_is_none(lang.modifier)

    def test_ll_cc_enc(self):
        lang = L.parse_language('el_GR.UTF-8')
        assert_equal(lang.language_code, 'el')
        assert_equal(lang.territory_code, 'GR')
        assert_equal(lang.encoding, 'UTF-8')
        assert_is_none(lang.modifier)

    def test_ll_cc_modifier(self):
        lang = L.parse_language('en_US@quot')
        assert_equal(lang.language_code, 'en')
        assert_equal(lang.territory_code, 'US')
        assert_is_none(lang.encoding)
        assert_equal(lang.modifier, 'quot')

    def test_syntax_error(self):
        with assert_raises(L.LanguageSyntaxError):
            L.parse_language('GR')

class test_get_primary_languages:

    def test_found(self):
        langs = L.get_primary_languages()
        assert_in('el', langs)

    def test_not_found(self):
        langs = L.get_primary_languages()
        assert_not_in('ry', langs)

    @collect_yielded
    def test_iso_639(self):
        def t(lang_str):
            lang = L.parse_language(lang_str)
            assert_is_none(lang.fix_codes())
            assert_equal(str(lang), lang_str)
        for lang_str in L.get_primary_languages():
            yield t, lang_str

class test_get_plural_forms:

    def t(self, lang):
        lang = L.parse_language(lang)
        return lang.get_plural_forms()

    def test_found_ll(self):
        assert_equal(
            self.t('el'),
            ['nplurals=2; plural=n != 1;']
        )

    def test_found_ll_cc(self):
        assert_equal(
            self.t('el_GR'),
            ['nplurals=2; plural=n != 1;']
        )

    def test_en_ca(self):
        assert_equal(
            self.t('en'),
            self.t('en_CA'),
        )

    def test_pt_br(self):
        assert_not_equal(
            self.t('pt'),
            self.t('pt_BR'),
        )

    def test_not_known(self):
        assert_is_none(self.t('la'))

    def test_not_found(self):
        assert_is_none(self.t('ry'))

class test_principal_territory:

    def test_found_2(self):
        # el -> el_GR
        lang = L.parse_language('el')
        cc = lang.get_principal_territory_code()
        assert_equal(cc, 'GR')

    def test_remove_2(self):
        # el_GR -> el
        lang = L.parse_language('el_GR')
        assert_equal(str(lang), 'el_GR')
        rc = lang.remove_principal_territory_code()
        assert_is(rc, True)
        assert_equal(str(lang), 'el')

    def test_found_3(self):
        # ang -> ang_GB
        lang = L.parse_language('ang')
        cc = lang.get_principal_territory_code()
        assert_equal(cc, 'GB')

    def test_remove_3(self):
        # ang_GB -> ang
        lang = L.parse_language('ang_GB')
        assert_equal(str(lang), 'ang_GB')
        rc = lang.remove_principal_territory_code()
        assert_is(rc, True)
        assert_equal(str(lang), 'ang')

    def test_no_principal_territory_code(self):
        # en -/-> en_US
        lang = L.parse_language('en')
        cc = lang.get_principal_territory_code()
        assert_is_none(cc)

    def test_no_remove_principal_territory_code(self):
        # en_US -/-> en
        lang = L.parse_language('en_US')
        assert_equal(str(lang), 'en_US')
        rc = lang.remove_principal_territory_code()
        assert_is_none(rc)
        assert_equal(str(lang), 'en_US')

    def test_not_found(self):
        lang = L.parse_language('ry')
        cc = lang.get_principal_territory_code()
        assert_is_none(cc)

class test_unrepresentable_characters:

    def test_ll_bad(self):
        lang = L.parse_language('pl')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_not_equal(result, [])

    def test_ll_ok(self):
        lang = L.parse_language('pl')
        result = lang.get_unrepresentable_characters('ISO-8859-2')
        assert_equal(result, [])

    def test_ll_cc_bad(self):
        lang = L.parse_language('pl_PL')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_not_equal(result, [])

    def test_ll_cc_ok(self):
        lang = L.parse_language('pl_PL')
        result = lang.get_unrepresentable_characters('ISO-8859-2')
        assert_equal(result, [])

    def test_ll_mod_bad(self):
        lang = L.parse_language('en@quot')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_not_equal(result, [])

    def test_ll_mod_ok(self):
        lang = L.parse_language('en@quot')
        result = lang.get_unrepresentable_characters('UTF-8')
        assert_equal(result, [])

    def test_ll_cc_mod_bad(self):
        lang = L.parse_language('en_US@quot')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_not_equal(result, [])

    def test_ll_cc_mod_ok(self):
        lang = L.parse_language('en_US@quot')
        result = lang.get_unrepresentable_characters('UTF-8')
        assert_equal(result, [])

    def test_ll_optional(self):
        # U+0178 (LATIN CAPITAL LETTER Y WITH DIAERESIS) is not representable
        # in ISO-8859-1, but we normally turn a blind eye to this.
        lang = L.parse_language('fr')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_equal(result, [])
        result = lang.get_unrepresentable_characters('ISO-8859-1', strict=True)
        assert_not_equal(result, [])

    def test_ll_not_found(self):
        lang = L.parse_language('ry')
        result = lang.get_unrepresentable_characters('ISO-8859-1')
        assert_is_none(result)

    @tools.fork_isolation
    def test_extra_encoding(self):
        encoding = 'GEORGIAN-PS'
        lang = L.parse_language('pl')
        with assert_raises(LookupError):
            ''.encode(encoding)
        E.install_extra_encodings()
        result = lang.get_unrepresentable_characters(encoding)
        assert_not_equal(result, [])

@collect_yielded
def test_glibc_supported():
    def t(l):
        lang = L.parse_language(l)
        try:
            lang.fix_codes()
        except L.FixingLanguageCodesFailed:
            # FIXME: some ISO-639-3 codes are not recognized yet
            if len(l.split('_')[0]) == 3:
                raise unittest.SkipTest('expected failure')
            reason = locales_to_skip.get(l)
            if reason is not None:
                raise unittest.SkipTest(reason)
            raise
        assert_equal(str(lang), l)
    try:
        file = open('/usr/share/i18n/SUPPORTED', encoding='ASCII')  # pylint: disable=consider-using-with
    except OSError as exc:
        raise unittest.SkipTest(exc)
    locales = set()
    with file:
        for line in file:
            if line[:1] in {'#', '\n'}:
                continue
            locale, *rest = line.split()
            del rest
            if (locale + '.').startswith('iw_IL.'):
                # iw_IL is obsolete
                continue
            if locale.startswith('C.'):
                continue
            locales.add(locale)
    misnamed_locales = {
        # glibc 2.21 had two misnamed locales:
        'bh_IN.UTF-8',  # should be bhb_IN
        'tu_IN.UTF-8',  # should be tcy_IN
    }
    locales_to_skip = {}
    if locales & misnamed_locales == misnamed_locales:
        for l in misnamed_locales:
            locales_to_skip[l] = 'https://sourceware.org/git/gitweb.cgi?p=glibc.git;a=commitdiff;h=032c510db06c'
    for l in sorted(locales):
        yield t, l

@collect_yielded
def test_poedit():
    # https://github.com/vslavik/poedit/blob/v1.8.1-oss/src/language_impl_legacy.h
    # There won't be any new names in this table,
    # so it's safe to hardcode them all here.
    def t(name, poedit_ll):
        poedit_ll = L.parse_language(poedit_ll)
        ll = L.get_language_for_name(name)
        assert_equal(ll, poedit_ll)
    def x(name, poedit_ll):
        poedit_ll = L.parse_language(poedit_ll)
        with assert_raises(LookupError):
            L.get_language_for_name(name)
        raise unittest.SkipTest('expected failure')
    yield t, 'Abkhazian', 'ab'
    yield t, 'Afar', 'aa'
    yield t, 'Afrikaans', 'af'
    yield t, 'Albanian', 'sq'
    yield t, 'Amharic', 'am'
    yield t, 'Arabic', 'ar'
    yield t, 'Armenian', 'hy'
    yield t, 'Assamese', 'as'
    yield t, 'Avestan', 'ae'
    yield t, 'Aymara', 'ay'
    yield t, 'Azerbaijani', 'az'
    yield t, 'Bashkir', 'ba'
    yield t, 'Basque', 'eu'
    yield t, 'Belarusian', 'be'
    yield t, 'Bengali', 'bn'
    yield t, 'Bislama', 'bi'
    yield t, 'Bosnian', 'bs'
    yield t, 'Breton', 'br'
    yield t, 'Bulgarian', 'bg'
    yield t, 'Burmese', 'my'
    yield t, 'Catalan', 'ca'
    yield t, 'Chamorro', 'ch'
    yield t, 'Chechen', 'ce'
    yield t, 'Chichewa; Nyanja', 'ny'
    yield t, 'Chinese', 'zh'
    yield t, 'Church Slavic', 'cu'
    yield t, 'Chuvash', 'cv'
    yield t, 'Cornish', 'kw'
    yield t, 'Corsican', 'co'
    yield t, 'Croatian', 'hr'
    yield t, 'Czech', 'cs'
    yield t, 'Danish', 'da'
    yield t, 'Dutch', 'nl'
    yield t, 'Dzongkha', 'dz'
    yield t, 'English', 'en'
    yield t, 'Esperanto', 'eo'
    yield t, 'Estonian', 'et'
    yield t, 'Faroese', 'fo'
    yield t, 'Fijian', 'fj'
    yield t, 'Finnish', 'fi'
    yield t, 'French', 'fr'
    yield t, 'Frisian', 'fy'
    yield t, 'Friulian', 'fur'
    yield t, 'Gaelic', 'gd'
    yield t, 'Galician', 'gl'
    yield t, 'Georgian', 'ka'
    yield t, 'German', 'de'
    yield t, 'Greek', 'el'
    yield t, 'Guarani', 'gn'
    yield t, 'Gujarati', 'gu'
    yield t, 'Hausa', 'ha'
    yield t, 'Hebrew', 'he'
    yield t, 'Herero', 'hz'
    yield t, 'Hindi', 'hi'
    yield t, 'Hiri Motu', 'ho'
    yield t, 'Hungarian', 'hu'
    yield t, 'Icelandic', 'is'
    yield t, 'Indonesian', 'id'
    yield t, 'Interlingua', 'ia'
    yield t, 'Interlingue', 'ie'
    yield t, 'Inuktitut', 'iu'
    yield t, 'Inupiaq', 'ik'
    yield t, 'Irish', 'ga'
    yield t, 'Italian', 'it'
    yield t, 'Japanese', 'ja'
    yield t, 'Javanese', 'jv'  # https://github.com/vslavik/poedit/pull/193
    yield t, 'Kalaallisut', 'kl'
    yield t, 'Kannada', 'kn'
    yield t, 'Kashmiri', 'ks'
    yield t, 'Kazakh', 'kk'
    yield t, 'Khmer', 'km'
    yield t, 'Kikuyu', 'ki'
    yield t, 'Kinyarwanda', 'rw'
    yield t, 'Komi', 'kv'
    yield t, 'Korean', 'ko'
    yield t, 'Kuanyama', 'kj'
    yield t, 'Kurdish', 'ku'
    yield t, 'Kyrgyz', 'ky'
    yield t, 'Lao', 'lo'
    yield t, 'Latin', 'la'
    yield t, 'Latvian', 'lv'
    yield t, 'Letzeburgesch', 'lb'
    yield t, 'Lingala', 'ln'
    yield t, 'Lithuanian', 'lt'
    yield t, 'Macedonian', 'mk'
    yield t, 'Malagasy', 'mg'
    yield t, 'Malay', 'ms'
    yield t, 'Malayalam', 'ml'
    yield t, 'Maltese', 'mt'
    yield t, 'Maori', 'mi'
    yield t, 'Marathi', 'mr'
    yield t, 'Marshall', 'mh'
    yield t, 'Moldavian', 'ro_MD'  # XXX poedit uses deprecated "mo"
    yield t, 'Mongolian', 'mn'
    yield t, 'Nauru', 'na'
    yield t, 'Navajo', 'nv'
    yield t, 'Ndebele, South', 'nr'
    yield t, 'Ndonga', 'ng'
    yield t, 'Nepali', 'ne'
    yield t, 'Northern Sami', 'se'
    yield t, 'Norwegian Bokmal', 'nb'
    yield t, 'Norwegian Nynorsk', 'nn'
    yield t, 'Occitan', 'oc'
    yield t, 'Oriya', 'or'
    yield t, 'Ossetian; Ossetic', 'os'
    yield t, 'Pali', 'pi'
    yield t, 'Panjabi', 'pa'
    yield t, 'Pashto, Pushto', 'ps'
    yield t, 'Persian', 'fa'
    yield t, 'Polish', 'pl'
    yield t, 'Portuguese', 'pt'
    yield t, 'Quechua', 'qu'
    yield t, 'Rhaeto-Romance', 'rm'
    yield t, 'Romanian', 'ro'
    yield t, 'Rundi', 'rn'
    yield t, 'Russian', 'ru'
    yield t, 'Samoan', 'sm'
    yield t, 'Sangro', 'sg'
    yield t, 'Sanskrit', 'sa'
    yield t, 'Sardinian', 'sc'
    yield t, 'Serbian', 'sr'
    yield t, 'Sesotho', 'st'
    yield t, 'Setswana', 'tn'
    yield t, 'Shona', 'sn'
    yield t, 'Sindhi', 'sd'
    yield t, 'Sinhalese', 'si'
    yield t, 'Siswati', 'ss'
    yield t, 'Slovak', 'sk'
    yield t, 'Slovenian', 'sl'
    yield t, 'Somali', 'so'
    yield t, 'Spanish', 'es'
    yield t, 'Sundanese', 'su'
    yield t, 'Swahili', 'sw'
    yield t, 'Swedish', 'sv'
    yield t, 'Tagalog', 'tl'
    yield t, 'Tahitian', 'ty'
    yield t, 'Tajik', 'tg'
    yield t, 'Tamil', 'ta'
    yield t, 'Tatar', 'tt'
    yield t, 'Telugu', 'te'
    yield t, 'Thai', 'th'
    yield t, 'Tibetan', 'bo'
    yield t, 'Tigrinya', 'ti'
    yield t, 'Tonga', 'to'
    yield t, 'Tsonga', 'ts'
    yield t, 'Turkish', 'tr'
    yield t, 'Turkmen', 'tk'
    yield t, 'Twi', 'tw'
    yield t, 'Ukrainian', 'uk'
    yield t, 'Urdu', 'ur'
    yield t, 'Uyghur', 'ug'
    yield t, 'Uzbek', 'uz'
    yield t, 'Vietnamese', 'vi'
    yield t, 'Volapuk', 'vo'
    yield t, 'Walloon', 'wa'
    yield t, 'Welsh', 'cy'
    yield t, 'Wolof', 'wo'
    yield t, 'Xhosa', 'xh'
    yield t, 'Yiddish', 'yi'
    yield t, 'Yoruba', 'yo'
    yield t, 'Zhuang', 'za'
    yield t, 'Zulu', 'zu'
    # TODO:
    yield x, '(Afan) Oromo', 'om'
    yield x, 'Bihari', 'bh'  # "bh" is marked as collective in ISO-639-2
    yield x, 'Serbian (Latin)', 'sr_RS@latin'
    yield x, 'Serbo-Croatian', 'sh'  # "sh" is deprecated in favor of "sr", "hr", "bs"

# vim:ts=4 sts=4 sw=4 et


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""dbus-sma.py: Driver to integrate SMA SunnyIsland inverters
                with Victron Venus OS. """

__author__      = "github usernames: madsci1016, jaedog"
__copyright__   = "Copyright 2020"
__license__     = "MIT"
__version__     = "1.1"

import os
import signal
import sys
import argparse
import serial
import socket
import logging
import yaml

from dbus.mainloop.glib import DBusGMainLoop
import dbus
import gobject

import can
from can.bus import BusState
from timeit import default_timer as timer
import time
from datetime import datetime, timedelta

# Victron packages
sys.path.insert(1, os.path.join(os.path.dirname(__file__), 'ext', 'velib_python'))
from vedbus import VeDbusService
from ve_utils import get_vrm_portal_id, exit_on_error
from dbusmonitor import DbusMonitor
from settingsdevice import SettingsDevice  # available in the velib_python repository

from bms_state_machine import BMSChargeStateMachine, BMSChargeModel, BMSChargeController

#from settingsdevice import SettingsDevice
#from logger import setup_logging
#import delegates
#from sc_utils import safeadd as _safeadd, safemax as _safemax

# ignore terminal resize signals (keeps exception from being thrown)
signal.signal(signal.SIGWINCH, signal.SIG_IGN)

softwareVersion = '1.1'
#logger = logging.getLogger("dbus-sma")
#logger = logging.getLogger(__name__)

# global logger for all modules imported here
logger = logging.getLogger()

#logging.basicConfig(filename='/data/etc/dbus-sma/logging.log', encoding='utf-8', level=logging.INFO)
#logger.setLevel(logging.DEBUG)
logger.setLevel(logging.INFO)

# The CANable (https://canable.io/) is a small open-source USB to CAN adapter. The CANable can show up as a virtual serial port (slcan): /dev/ttyACM0 or
# as a socketcan: can0. In testing both methods work, however, I found the can0 to be much more robust.
# Devices from http://protofusion.org store ship by default with the "slcan" firmware. It can be flashed with the "candlelight" firmware to
# support socketcan.

# When the adapter is a socketcan, bring up link first as root:
# ip link set can0 up type can bitrate 500000
#
# To capture CAN msgs on the bus:
# tcpdump -w capture.pcap -i can5

# TODO: change to input param
#canBusChannel = "/dev/ttyACM0"
canBusChannel = "can5"

#canBusType = "slcan"
canBusType = "socketcan"

# connect and register to dbus
driver = {
	'name'        : "SMA SunnyIsland",
	'servicename' : "smasunnyisland",
	'instance'    : 261,
	'id'          : 2754,
	'version'     : 476,
	'serial'      : "SMABillConnect",
	'connection'  : "com.victronenergy.vebus.smasunnyisland"
}

CAN_tx_msg = {"BatChg": 0x351, "BatSoC": 0x355, "BatVoltageCurrent" : 0x356, "AlarmWarning": 0x35a, "BMSOem": 0x35e, "BatData": 0x35f}
CANFrames = {"ExtPwr": 0x300, "InvPwr": 0x301, "OutputVoltage": 0x304, "Battery": 0x305, "Relay": 0x306, "Bits": 0x307, "LoadPwr": 0x308, "ExtVoltage": 0x309}
sma_line1 = {"OutputVoltage": 0, "ExtPwr": 0, "InvPwr": 0, "ExtVoltage": 0, "ExtFreq": 0.00, "OutputFreq": 0.00}
sma_line2 = {"OutputVoltage": 0, "ExtPwr": 0, "InvPwr": 0, "ExtVoltage": 0}
sma_battery = {"Voltage": 0, "Current": 0}
sma_system = {"State": 0, "ExtRelay" : 0, "ExtOk" : 0, "Load" : 0}

settings = 0

#command packets to turn SMAs on or off
SMA_ON_MSG = can.Message(arbitration_id = 0x35C,    #on
      data=[0b00000001,0,0,0],
      is_extended_id=False)

SMA_OFF_MSG = can.Message(arbitration_id = 0x35C,    #off
      data=[0b00000010,0,0,0],
      is_extended_id=False)

def getSignedNumber(number, bitLength):
    mask = (2 ** bitLength) - 1
    if number & (1 << (bitLength - 1)):
        return number | ~mask
    else:
        return number & mask

def bytes(integer):
    return divmod(integer, 0x100)

class BMSData:
  def __init__(self, max_battery_voltage, min_battery_voltage, low_battery_voltage, \
    charge_bulk_amps, max_discharge_amps, charge_absorb_voltage, charge_float_voltage, \
    time_min_absorb, rebulk_voltage):

    # settings for BMS

    # max and min battery voltage is used by the SMA as fault values
    # if the voltage goes above max or below min, the SMA will fault OFF
    # the inverter.
    self.max_battery_voltage = max_battery_voltage
    self.min_battery_voltage = min_battery_voltage

    # low battery voltage is used to trigger the SMA to connect to grid and
    # begin charging the batteries. Note, this value must be greater than
    # the min_battery_voltage
    self.low_battery_voltage = low_battery_voltage

    self.charge_bulk_amps = charge_bulk_amps
    self.max_discharge_amps = max_discharge_amps
    self.charge_absorb_voltage = charge_absorb_voltage
    self.charge_float_voltage = charge_float_voltage
    self.time_min_absorb = time_min_absorb
    self.rebulk_voltage = rebulk_voltage

    # state of BMS
    self.charging_state = "" # state of charge state machine
    self.state_of_charge = 42.0  # sane initial value
    self.actual_battery_voltage = 0.0
    self.req_discharge_amps = max_discharge_amps
    self.battery_current = 0.0
    self.pv_current = 0.0

  def __str__(self):
    return "BMS Data, MaxV: {0}V, MinV: {1}V, LowV: {2}V, BulkA: {3}A, AbsorbV: {4}V, FloatV: {5}V, MinuteAbsorb: {6}, RebulkV: {7}V" \
      .format(self.max_battery_voltage, self.min_battery_voltage, self.low_battery_voltage, self.charge_bulk_amps, \
        self.charge_absorb_voltage, self.charge_float_voltage, self.time_min_absorb, self.rebulk_voltage)

# SMA Driver Class
class SmaDriver:

  def __init__(self):
    self.driver_start_time = datetime.now()

    # data from yaml config file
    self._cfg = self.get_config_data()
    _cfg_bms = self._cfg['BMSData']

    # TODO: use venus settings to define these values
    #Initial BMS values eventually read from settings.
    self._bms_data = BMSData(max_battery_voltage=_cfg_bms['max_battery_voltage'], \
      min_battery_voltage=_cfg_bms['min_battery_voltage'], low_battery_voltage=_cfg_bms['low_battery_voltage'], \
      charge_bulk_amps=_cfg_bms['charge_bulk_amps'], max_discharge_amps=_cfg_bms['max_discharge_amps'], \
      charge_absorb_voltage=_cfg_bms['charge_absorb_voltage'], charge_float_voltage=_cfg_bms['charge_float_voltage'], \
      time_min_absorb=_cfg_bms['time_min_absorb'], rebulk_voltage=_cfg_bms['rebulk_voltage'])

    self.bms_controller = BMSChargeController(charge_bulk_current=self._bms_data.charge_bulk_amps, \
      charge_absorb_voltage=self._bms_data.charge_absorb_voltage, charge_float_voltage=self._bms_data.charge_float_voltage, \
        time_min_absorb=self._bms_data.time_min_absorb, rebulk_voltage=self._bms_data.rebulk_voltage)
    ret = self.bms_controller.start_charging()

    # Have a mainloop, so we can send/receive asynchronous calls to and from dbus
    DBusGMainLoop(set_as_default=True)

    self._can_bus = False

    self._safety_off = False   #flag to see if we every shut the inverters off due to low batt.

    logger.debug("Can bus init")
    try :
      self._can_bus = can.interface.Bus(bustype=canBusType, channel=canBusChannel, bitrate=500000)
    except can.CanError as e:
     logger.error(e)

    logger.debug("Can bus init done")

    # Add the AcInput1 setting if it doesn't exist so that the grid data is reported
    # to the system by dbus-systemcalc-py service
    settings = SettingsDevice(
       bus=dbus.SystemBus(),# if (platform.machine() == 'armv7l') else dbus.SessionBus(),
       supportedSettings={
           'acinput': ['/Settings/SystemSetup/AcInput1', 1, 0, 0],
           'hub4mode': ['/Settings/CGwacs/Hub4Mode', 3, 0, 0],
           'gridmeter': ['/Settings/CGwacs/RunWithoutGridMeter', 1, 0, 0],
           'acsetpoint': ['/Settings/CGwacs/AcPowerSetPoint', 0, 0, 0],
           'maxchargepwr': ['/Settings/CGwacs/MaxChargePower', 0, 0, 0],
           'maxdischargepwr': ['/Settings/CGwacs/MaxDischargePower', 0, 0, 0],
           'maxchargepercent': ['/Settings/CGwacs/MaxChargePercentage', 0, 0, 0],
           'maxdischargepercent': ['/Settings/CGwacs/MaxDischargePercentage', 0, 0, 0],
           'essMode': ['/Settings/CGwacs/BatteryLife/State', 0, 0, 0],
           },
       eventCallback=None)

		# Why this dummy? Because DbusMonitor expects these values to be there, even though we don't
		# need them. So just add some dummy data. This can go away when DbusMonitor is more generic.
    dummy = {'code': None, 'whenToLog': 'configChange', 'accessLevel': None}
    dbus_tree = {'com.victronenergy.system':
      {'/Dc/Battery/Soc': dummy, '/Dc/Battery/Current': dummy, '/Dc/Battery/Voltage': dummy, \
        '/Dc/Pv/Current': dummy, '/Ac/PvOnOutput/L1/Power': dummy, '/Ac/PvOnOutput/L2/Power': dummy, }}

    self._dbusmonitor = self._create_dbus_monitor(dbus_tree, valueChangedCallback=self._dbus_value_changed)

    self._dbusservice = self._create_dbus_service()

    self._dbusservice.add_path('/Serial',        value=12345)

    # /SystemState/State   ->   0: Off
    #                      ->   1: Low power
    #                      ->   2: VE.Bus Fault condition
    #                      ->   3: Bulk charging
    #                      ->   4: Absorption charging
    #                      ->   5: Float charging
    #                      ->   6: Storage mode
    #                      ->   7: Equalisation charging
    #                      ->   8: Passthru
    #                      ->   9: Inverting
    #                      ->  10: Assisting
    #                      -> 256: Discharging
    #                      -> 257: Sustain
    self._dbusservice.add_path('/State',                   0)
    self._dbusservice.add_path('/Mode',                    3)
    self._dbusservice.add_path('/Ac/PowerMeasurementType', 0)
    self._dbusservice.add_path('/Hub4/AssistantId', 5)
    self._dbusservice.add_path('/Hub4/DisableCharge', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/DisableFeedIn', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/DoNotFeedInOverVoltage', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/L1/AcPowerSetpoint', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/L2/AcPowerSetpoint', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/Sustain', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/L1/MaxFeedInPower', value=0, writeable=True)
    self._dbusservice.add_path('/Hub4/L2/MaxFeedInPower', value=0, writeable=True)

    # Create the inverter/charger paths
    self._dbusservice.add_path('/Ac/Out/L1/P',            -1)
    self._dbusservice.add_path('/Ac/Out/L2/P',            -1)
    self._dbusservice.add_path('/Ac/Out/L1/I',            -1)
    self._dbusservice.add_path('/Ac/Out/L2/I',            -1)
    self._dbusservice.add_path('/Ac/Out/L1/V',            -1)
    self._dbusservice.add_path('/Ac/Out/L2/V',            -1)
    self._dbusservice.add_path('/Ac/Out/L1/F',            -1)
    self._dbusservice.add_path('/Ac/Out/L2/F',            -1)
    self._dbusservice.add_path('/Ac/Out/P',               -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L1/P',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L2/P',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/P',          -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L1/V',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L2/V',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L1/F',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L2/F',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L1/I',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/L2/I',       -1)
    self._dbusservice.add_path('/Ac/ActiveIn/Connected',   1)
    self._dbusservice.add_path('/Ac/ActiveIn/ActiveInput', 0)
    self._dbusservice.add_path('/VebusError',              0)
    self._dbusservice.add_path('/Dc/0/Voltage',           -1)
    self._dbusservice.add_path('/Dc/0/Power',             -1)
    self._dbusservice.add_path('/Dc/0/Current',           -1)
    self._dbusservice.add_path('/Ac/NumberOfPhases',       2)
    self._dbusservice.add_path('/Alarms/GridLost',         0)

    # /VebusChargeState  <- 1. Bulk
    #                       2. Absorption
    #                       3. Float
    #                       4. Storage
    #                       5. Repeat absorption
    #                       6. Forced absorption
    #                       7. Equalise
    #                       8. Bulk stopped
    self._dbusservice.add_path('/VebusChargeState',        0)

    # Some attempts at logging consumption. Float of kwhr since driver start (i think)
    self._dbusservice.add_path('/Energy/GridToDc',         0)
    self._dbusservice.add_path('/Energy/GridToAcOut',      0)
    self._dbusservice.add_path('/Energy/DcToAcOut',        0)
    self._dbusservice.add_path('/Energy/AcIn1ToInverter',  0)
    self._dbusservice.add_path('/Energy/AcIn1ToAcOut',     0)
    self._dbusservice.add_path('/Energy/InverterToAcOut',  0)
    self._dbusservice.add_path('/Energy/Time',       timer())

    self._changed = True

    # create timers (time in msec)
    gobject.timeout_add(2000, exit_on_error, self._can_bus_txmit_handler)
    gobject.timeout_add(2000, exit_on_error, self._energy_handler)
    gobject.timeout_add(20, exit_on_error, self._parse_can_data_handler)

#----
  def __del__(self):
    if (self._can_bus):
      self._can_bus.shutdown()
      self._can_bus = False
      logger.debug("bus shutdown")

#----
  def run(self):
    # Start and run the mainloop
    logger.info("Starting mainloop, responding only on events")
    self._mainloop = gobject.MainLoop()

    try:
      self._mainloop.run()
    except KeyboardInterrupt:
      self._mainloop.quit()

#----
  def _create_dbus_monitor(self, *args, **kwargs):
    return DbusMonitor(*args, **kwargs)

#----
  def _create_dbus_service(self):
    dbusservice = VeDbusService(driver['connection'])
    dbusservice.add_mandatory_paths(
      processname=__file__,
      processversion=softwareVersion,
      connection=driver['connection'],
      deviceinstance=driver['instance'],
      productid=driver['id'],
      productname=driver['name'],
      firmwareversion=driver['version'],
      hardwareversion=driver['version'],
      connected=1)
    return dbusservice

#----
  # callback that gets called ever time a dbus value has changed
  def _dbus_value_changed(self, dbusServiceName, dbusPath, dict, changes, deviceInstance):
    self._changed = True

#----
  # called by timer every 20 msec
  def _parse_can_data_handler(self):

    try:
      msg = None
      # read msgs until we get one we want
      while True:
        msg = self._can_bus.recv(1)
        if (msg is None) :
          sma_system["State"] = 0
          #self._dbusservice["/State"] = 0
          logger.info("No Message received from Sunny Island")
          return True

        if (msg.arbitration_id == CANFrames["ExtPwr"] or msg.arbitration_id == CANFrames["InvPwr"] or \
              msg.arbitration_id == CANFrames["LoadPwr"] or msg.arbitration_id == CANFrames["OutputVoltage"] or \
              msg.arbitration_id == CANFrames["ExtVoltage"] or msg.arbitration_id == CANFrames["Battery"] or \
              msg.arbitration_id == CANFrames["Relay"] or msg.arbitration_id == CANFrames["Bits"]):
          break

      if msg is not None:
        if msg.arbitration_id == CANFrames["ExtPwr"]:
          sma_line1["ExtPwr"] = (getSignedNumber(msg.data[0] + msg.data[1]*256, 16)*100)
          sma_line2["ExtPwr"] = (getSignedNumber(msg.data[2] + msg.data[3]*256, 16)*100)
          #self._updatedbus()
          #print ("Ex Power L1: " + str(sma_line1["ExtPwr"]) + "  Power L2: " + str(sma_line2["ExtPwr"]))
        elif msg.arbitration_id == CANFrames["InvPwr"]:
          sma_line1["InvPwr"] = (getSignedNumber(msg.data[0] + msg.data[1]*256, 16)*100)
          sma_line2["InvPwr"] = (getSignedNumber(msg.data[2] + msg.data[3]*256, 16)*100)
          #calculate_pwr()
          #print ("Power L1: " + str(sma_line1["InvPwr"]) + "  Power L2: " + str(sma_line2["InvPwr"]))
          self._updatedbus()
        elif msg.arbitration_id == CANFrames["LoadPwr"]:
          sma_system["Load"] = (getSignedNumber(msg.data[0] + msg.data[1]*256, 16)*100)
          self._updatedbus()
        elif msg.arbitration_id == CANFrames["OutputVoltage"]:
          sma_line1["OutputVoltage"] = (float(getSignedNumber(msg.data[0] + msg.data[1]*256, 16))/10)
          sma_line2["OutputVoltage"] = (float(getSignedNumber(msg.data[2] + msg.data[3]*256, 16))/10)
          sma_line1["OutputFreq"] = float(msg.data[6] + msg.data[7]*256) / 100
          self._updatedbus()
        elif msg.arbitration_id == CANFrames["ExtVoltage"]:
          sma_line1["ExtVoltage"] = (float(getSignedNumber(msg.data[0] + msg.data[1]*256, 16))/10)
          sma_line2["ExtVoltage"] = (float(getSignedNumber(msg.data[2] + msg.data[3]*256, 16))/10)
          sma_line1["ExtFreq"] = float(msg.data[6] + msg.data[7]*256) / 100
          self._updatedbus()
        elif msg.arbitration_id == CANFrames["Battery"]:
          sma_battery["Voltage"] = float(msg.data[0] + msg.data[1]*256) / 10
          sma_battery["Current"] = float(getSignedNumber(msg.data[2] + msg.data[3]*256, 16)) / 10
          self._updatedbus()
        elif msg.arbitration_id == CANFrames["Bits"]:
          if msg.data[2]&128:
            sma_system["ExtRelay"] = 1
          else:
            sma_system["ExtRelay"] = 0
          if msg.data[2]&64:
            sma_system["ExtOk"] = 0
            #print ("Grid OK")
          else:
            #it seems to always report grid down once during relay transfer, so lets wait for two messages to latch.
            if sma_system["ExtOk"] == 0:
              sma_system["ExtOk"] = 1
            elif sma_system["ExtOk"] == 1:
              sma_system["ExtOk"] = 2
            #print ("Grid Down")

          #print ("307 message" )
          #print(msg)

    except (KeyboardInterrupt) as e:
      self._mainloop.quit()
    except (can.CanError) as e:
      logger.error(e)
      pass
    except Exception as e:
      exception_type = type(e).__name__
      logger.error("Exception occured: {0}, {1}".format(exception_type, e))

    return True

#----
  def _updatedbus(self):
    #self._dbusservice["/State"] = sma_system["State"]
    self._dbusservice["/Ac/ActiveIn/L1/P"] = sma_line1["ExtPwr"]
    self._dbusservice["/Ac/ActiveIn/L2/P"] = sma_line2["ExtPwr"]
    self._dbusservice["/Ac/ActiveIn/L1/V"] = sma_line1["ExtVoltage"]
    self._dbusservice["/Ac/ActiveIn/L2/V"] = sma_line2["ExtVoltage"]
    self._dbusservice["/Ac/ActiveIn/L1/F"] = sma_line1["ExtFreq"]
    self._dbusservice["/Ac/ActiveIn/L2/F"] = sma_line1["ExtFreq"]
    if sma_system["ExtOk"] == 0 or sma_system["ExtOk"] == 2:
      self._dbusservice["/Alarms/GridLost"] = sma_system["ExtOk"]
    if sma_line1["ExtVoltage"] != 0:
      self._dbusservice["/Ac/ActiveIn/L1/I"] = int(sma_line1["ExtPwr"] / sma_line1["ExtVoltage"])
    if sma_line2["ExtVoltage"] != 0:
      self._dbusservice["/Ac/ActiveIn/L2/I"] = int(sma_line2["ExtPwr"] / sma_line2["ExtVoltage"])
    self._dbusservice["/Ac/ActiveIn/P"] = sma_line1["ExtPwr"] + sma_line2["ExtPwr"]
    self._dbusservice["/Dc/0/Voltage"] = sma_battery["Voltage"]
    self._dbusservice["/Dc/0/Current"] = sma_battery["Current"] *-1
    self._dbusservice["/Dc/0/Power"] = sma_battery["Current"] * sma_battery["Voltage"] *-1

    line1_inv_outpwr = sma_line1["ExtPwr"] + sma_line1["InvPwr"]
    line2_inv_outpwr = sma_line2["ExtPwr"] + sma_line2["InvPwr"]

    #print ("After calc Power L1: " + str(line1_inv_outpwr) + "  Power L2: " + str(line2_inv_outpwr))

    #we can gain back a little bit of resolution by compairing total reported load to sum of line loads reported to remove one source of rounding error.
    if (sma_system["Load"] == (line1_inv_outpwr + line2_inv_outpwr + 100)):
      line1_inv_outpwr+=50
      line2_inv_outpwr+=50
    elif (sma_system["Load"] == (line1_inv_outpwr + line2_inv_outpwr - 100)):
      line1_inv_outpwr-=50
      line2_inv_outpwr-=50

    self._dbusservice["/Ac/Out/L1/P"] = line1_inv_outpwr
    self._dbusservice["/Ac/Out/L2/P"] = line2_inv_outpwr
    self._dbusservice["/Ac/Out/P"] =  sma_system["Load"]
    self._dbusservice["/Ac/Out/L1/F"] = sma_line1["OutputFreq"]
    self._dbusservice["/Ac/Out/L2/F"] = sma_line1["OutputFreq"]
    self._dbusservice["/Ac/Out/L1/V"] = sma_line1["OutputVoltage"]
    self._dbusservice["/Ac/Out/L2/V"] = sma_line2["OutputVoltage"]

    inverter_on = 0
    if sma_line1["OutputVoltage"] > 5:
      self._dbusservice["/Ac/Out/L1/I"] = int(line1_inv_outpwr / sma_line1["OutputVoltage"])
      inverter_on += 1
    if sma_line2["OutputVoltage"] > 5:
      self._dbusservice["/Ac/Out/L2/I"] = int(line2_inv_outpwr / sma_line2["OutputVoltage"])
      inverter_on += 1

    if sma_system["ExtRelay"]:
      self._dbusservice["/Ac/ActiveIn/Connected"] = 1
      self._dbusservice["/Ac/ActiveIn/ActiveInput"] = 0
    else:
      self._dbusservice["/Ac/ActiveIn/Connected"] = 0
      self._dbusservice["/Ac/ActiveIn/ActiveInput"] = 240

    # state = 3:Bulk, 4:Absorb, 5:Float, 6:Storage, 7:Equalize, 8:Passthrough 9:Inverting
    # push charging state to dbus
    vebusChargeState = 0
    sma_system["State"] = 0

    #logger.info("SysState: {0}, InvOn: {1}".format(systemState, inverter_on))

    if (inverter_on > 0):
      sma_system["State"] = 9
      # if current is going into the battery
      if (self._bms_data.battery_current > 0):
        if (self._bms_data.charging_state == "bulk_chg"):
          vebusChargeState = 1
          sma_system["State"] = 3
        elif (self._bms_data.charging_state == "absorb_chg"):
          vebusChargeState = 2
          sma_system["State"] = 4
        elif (self._bms_data.charging_state == "float_chg"):
          vebusChargeState = 3
          sma_system["State"] = 5

    self._dbusservice["/VebusChargeState"] = vebusChargeState
    self._dbusservice["/State"] = sma_system["State"]

#----
  def _energy_handler(self):
    energy_sec = timer() - self._dbusservice["/Energy/Time"]
    self._dbusservice["/Energy/Time"] = timer()

    if self._dbusservice["/Dc/0/Power"] > 0:
      #Grid to battery
      self._dbusservice["/Energy/GridToAcOut"] = self._dbusservice["/Energy/GridToAcOut"] + \
        ((self._dbusservice["/Ac/Out/P"]) * energy_sec * 0.00000028)

      self._dbusservice["/Energy/GridToDc"] = self._dbusservice["/Energy/GridToDc"] + \
        (self._dbusservice["/Dc/0/Power"]  * energy_sec * 0.00000028)
    else:
      #battery to out
      self._dbusservice["/Energy/DcToAcOut"] = self._dbusservice["/Energy/DcToAcOut"] + \
        ((self._dbusservice["/Ac/Out/P"])  * energy_sec * 0.00000028)

    #print(timer() - self._dbusservice["/Energy/Time"], ":", self._dbusservice["/Ac/Out/P"])

    self._dbusservice["/Energy/AcIn1ToAcOut"] = self._dbusservice["/Energy/GridToAcOut"]
    self._dbusservice["/Energy/AcIn1ToInverter"] = self._dbusservice["/Energy/GridToDc"]
    self._dbusservice["/Energy/InverterToAcOut"] = self._dbusservice["/Energy/DcToAcOut"]
    self._dbusservice["/Energy/Time"] = timer()
    return True

#----
  # BMS charge logic since SMA is in dumb mode
  def _execute_grid_solar_charge_logic(self):
    charge_amps = None

    # time in UTC
    now = datetime.now()

    # SMA Sunny Island Feature:
    # Setting 232# Grid Control
    # Item 41 GdSocEna - Activate the grid request based on SOC (Default: Disable) = Enable
    #
    # By enabling this setting the SMA will activate grid to charge batteries. To set the ranges:
    # Setting 233# Grid Start
    #
    # Item 01 GdSocTm1Str - SOC limit for switching on utility grid for time 1 = 40%
    # Item 02 GdSocTm1Stp - SOC limit for switching off the utility grid for time 1 = 80%
    # Item 03 GdSocTm2Str - SOC limit for switching on utility grid for time 2 = 40%
    # Item 04 GdSocTm2Stp - SOC limit for switching off the utility grid for time 2 = 80%

    if (sma_system["ExtRelay"] == 1):
      #no point in running the math below to calculate a new target charge current unless we have an update from the inverters
      #which is slow. Like every 12 seconds.
      #global SMAupdate
      #if SMAupdate == True:
      #  SMAupdate = False

      _cfg_grid = self._cfg["GridLogic"]
      _cfg_safety = self._cfg["SafetyLogic"]
      #requested charge current varies by time of day and SoC value
      #for now, some rules to change charge behavior hard coded for my application.
      #Gonna try making these charge current targets inlcuding solar, so we need to subtract solar current later.
      if now.hour >= _cfg_grid["start_hour"] and now.hour <= _cfg_grid["end_hour"]:
        if now.hour >= _cfg_grid["mid_hour"] and self._bms_data.state_of_charge < 49.0:
          charge_amps = _cfg_grid["mid_hour_current"]
        else:
          charge_amps = _cfg_grid["current"]
      else:
        charge_amps = _cfg_grid["offtime_current"]

      #TODO: can this use the same value as default bulk current?
      if self._bms_data.state_of_charge < _cfg_safety["after_blackout_min_soc"]:  #recovering from blackout? Charge fast!
        charge_amps = _cfg_safety["after_blackout_charge_amps"]

      #subtract any active Solar current from the requested charge current
      charge_amps = charge_amps - self._bms_data.pv_current

      # if pv_current is greater than requested charge amps, don't go negative
      if (charge_amps < 0.0):
        charge_amps = 0.0

    logger.info("Grid Logic: Time: {0}, On Grid: {1} Charge amps: {2}" \
      .format(now, sma_system["ExtRelay"], charge_amps))

    return charge_amps

#----
 	# Called on a two second timer to send CAN messages
  def _can_bus_txmit_handler(self):

    # log data received from SMA on CAN bus (doing it here since this timer is slower!)
    out_load_msg = "SMA: System Load: {0}, Driver runtime: {1}".format(sma_system["Load"], datetime.now() - self.driver_start_time)

    out_ext_msg = "SMA: External, Line 1: {0}V, Line 2: {1}V, Line 1 Pwr: {2}W, Line 2 Pwr: {3}W, Freq: {4}" \
      .format(sma_line1["ExtVoltage"], sma_line2["ExtVoltage"], sma_line1["ExtPwr"], sma_line2["ExtPwr"], sma_line1["ExtFreq"])

    out_inv_msg = "SMA: Inverter, Line 1: {0}V, Line 2: {1}V, Line 1 Pwr: {2}W, Line 2 Pwr: {3}W, Freq: {4}" \
      .format(sma_line1["OutputVoltage"], sma_line2["OutputVoltage"], sma_line1["InvPwr"], sma_line2["InvPwr"], sma_line1["OutputFreq"])

    out_batt_msg = "SMA: Batt Voltage: {0}, Batt Current: {1}" \
      .format(sma_battery["Voltage"], sma_battery["Current"])

    logger.info(out_load_msg)
    logger.info(out_ext_msg)
    logger.info(out_inv_msg)
    logger.info(out_batt_msg)

    #get some data from the Victron BUS, invalid data returns NoneType
    soc = self._dbusmonitor.get_value('com.victronenergy.system', '/Dc/Battery/Soc')
    volt = self._dbusmonitor.get_value('com.victronenergy.system', '/Dc/Battery/Voltage')
    current = self._dbusmonitor.get_value('com.victronenergy.system', '/Dc/Battery/Current')
    pv_current = self._dbusmonitor.get_value('com.victronenergy.system', '/Dc/Pv/Current')
    if (pv_current == None):
      pv_current = 0.0

    # if we don't have these values, there is nothing to do!
    if (soc == None or volt == None):
      logger.error("DBusMonitor returning None for one or more: SOC: {0}, Volt: {1}, Current: {2}, PVCurrent: {3}" \
          .format(soc, volt, current, pv_current))
      return True

    # update bms state data
    self._bms_data.state_of_charge = soc
    self._bms_data.actual_battery_voltage = volt
    self._bms_data.battery_current = current
    self._bms_data.pv_current = pv_current

    # update the requested bulk current based on the grid solar charge logic
    self.bms_controller.update_req_bulk_current(self._execute_grid_solar_charge_logic())

    # update the battery voltage for the BMS to determine next state or charge current level
    # Note: Positive value for current means it is going INTO the battery. SMA will report as negative
    # so we change signs here
    is_state_changed = self.bms_controller.update_battery_data(self._bms_data.actual_battery_voltage, \
        -(sma_battery["Current"]))

    self._bms_data.charging_state = self.bms_controller.get_state()
    charge_current = self.bms_controller.get_charge_current()

    logger.info("BMS Send, SoC: {0:.1f}%, Batt Voltage: {1:.2f}V, Batt Current: {2:.2f}A, Charge State: {3}, Req Charge: {4}A, Req Discharge: {5}A, PV Cur: {6} ". \
        format(self._bms_data.state_of_charge, self._bms_data.actual_battery_voltage, \
        self._bms_data.battery_current, self._bms_data.charging_state, charge_current,
        self._bms_data.req_discharge_amps, self._bms_data.pv_current))

    #**************Low battery safety****************#

    _cfg_safety = self._cfg["SafetyLogic"]

    #if grid is up but battery low voltage, issue with shunt calibration or SMA setting, pre-empt SoC with minimum value to force grid transfer
    if (sma_system["ExtOk"] == 0 and self._bms_data.actual_battery_voltage < self._bms_data.low_battery_voltage):
      self._bms_data.state_of_charge = 1.0

    #if no grid and Soc is low, we are in blackout with dead batteries and need to shut off inverters
    if(self._safety_off == False):
      #normal running, check for grid not ok AND low Soc, send off message till inverters respond
      if(sma_system["ExtOk"] == 2 and  soc < _cfg_safety["min_soc_inv_off"]):
        self._can_bus.send(SMA_OFF_MSG)
        if(sma_system["State"] == 0):
          self._safety_off = True
        #print("Shut off due to low SoC")
    else:
      #if we saftey shutdown, keep checking for grid restore OR SoC increase, send on message till inverters respond
      if(sma_system["ExtOk"] == 0 or soc >= _cfg_safety["min_soc_inv_off"]):
        self._can_bus.send(SMA_ON_MSG)
        if(sma_system["State"] != 0):
          self._safety_off = False
        #print("Start SMA due to grid restore or SoC increase")

    #breakup some of the values for CAN packing
    SoC_HD = int(self._bms_data.state_of_charge*100)
    SoC_HD_H, SoC_HD_L = bytes(SoC_HD)

    Req_Charge_H, Req_Charge_L = bytes(int(charge_current*10))

    Req_Discharge_H, Req_Discharge_L = bytes(int(self._bms_data.req_discharge_amps*10))
    Max_V_H, Max_V_L = bytes(int(self._bms_data.max_battery_voltage*10))
    Min_V_H, Min_V_L = bytes(int(self._bms_data.min_battery_voltage*10))

    msg = can.Message(arbitration_id = CAN_tx_msg["BatChg"],
      data=[Max_V_L, Max_V_H, Req_Charge_L, Req_Charge_H, Req_Discharge_L, Req_Discharge_H, Min_V_L, Min_V_H],
      is_extended_id=False)

    msg2 = can.Message(arbitration_id = CAN_tx_msg["BatSoC"],
      data=[int(self._bms_data.state_of_charge), 0x00, 0x64, 0x0, SoC_HD_L, SoC_HD_H],
      is_extended_id=False)

    msg3 = can.Message(arbitration_id = CAN_tx_msg["BatVoltageCurrent"],
      data=[0x00, 0x00, 0x00, 0x0, 0xf0, 0x00],
      is_extended_id=False)

    msg4 = can.Message(arbitration_id = CAN_tx_msg["AlarmWarning"],
      data=[0x00, 0x00, 0x00, 0x0, 0x00, 0x00, 0x00, 0x00],
      is_extended_id=False)

    msg5 = can.Message(arbitration_id = CAN_tx_msg["BMSOem"],
      data=[0x42, 0x41, 0x54, 0x52, 0x49, 0x55, 0x4d, 0x20],
      is_extended_id=False)

    msg6 = can.Message(arbitration_id = CAN_tx_msg["BatData"],
      data=[0x03, 0x04, 0x0a, 0x04, 0x76, 0x02, 0x00, 0x00],
      is_extended_id=False)

    #logger.debug(self._can_bus)

    try :
      self._can_bus.send(msg)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))
      time.sleep(.100)

      self._can_bus.send(msg2)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))
      time.sleep(.100)

      self._can_bus.send(msg3)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))

      time.sleep(.100)

      self._can_bus.send(msg4)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))

      time.sleep(.100)

      self._can_bus.send(msg5)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))

      time.sleep(.100)

      self._can_bus.send(msg6)
      #logger.debug("Message sent on {}".format(self._can_bus.channel_info))

      #logger.info("Sent to SI: {0}, {1}, {2}, {3}, {4}". \
      #  format(self._bms_data.req_discharge_amps, self._bms_data.state_of_charge, \
      #  self._bms_data.actual_battery_voltage, self._bms_data.battery_current, \
      #  self._bms_data.pv_current))

      #logger.info("Sent 6 messages on {}".format(self._can_bus.channel_info))
    except (can.CanError) as e:
      logger.error("CAN BUS Transmit error (is controller missing?): %s" % e.message)
    except KeyboardInterrupt:
      pass

    return True  # keep timer running

#----
  def get_config_data(self):
    try :
      dir_path = os.path.dirname(os.path.realpath(__file__))
      with open(dir_path + "/dbus-sma.yaml", "r") as yamlfile:
        config = yaml.load(yamlfile, Loader=yaml.FullLoader)
        return config
    except :
      logger.info("dbus-sma.yaml file not found or correct.")
      sys.exit()

if __name__ == "__main__":
  # Argument parsing
  parser = argparse.ArgumentParser(description='Converts readings from AC-Sensors connected to a VE.Bus device in a pvinverter ' + 'D-Bus service.')
  parser.add_argument('-s', '--serial', help='tty')
  parser.add_argument("-d", "--debug", help="set logging level to debug",action="store_true")

  args = parser.parse_args()

  print("-------- dbus_SMADriver, v" + softwareVersion + " is starting up --------")
  #logger = setup_logging(args.debug)

  # create SMA Driver
  smadriver = SmaDriver()

  # run driver (starts mainloop and hangs until CTRL+C/SIGINT received)
  smadriver.run()

  # force clean up resources
  smadriver.__del__()

  print("-------- dbus_SMADriver, v" + softwareVersion + " is shuting down --------")

  sys.exit(1)


# coding=utf-8
# Se añade esta linea para poner tildes

from django import forms
from .models import Post

class PostForm(forms.ModelForm):
  # Nombre del label del imput. La variable es la misma que se definio en el modelo
  #title = forms.CharField(label='Título')

  ''' Esto se hace para dar atributos a los campos del formulario
  def __init__(self, *args, **kwargs):
    super(PostForm, self).__init__(*args, **kwargs)
    self.fields['title'].widget.attrs = {
      'title': 'Titulo'
    }
  '''

  class Meta:
    model = Post
    fields = [
      'title',
      'text'
    ]


import csv
from sqlalchemy.orm import sessionmaker
from .create_tables import BillOfMaterials
from etl_scripts.config import oltp_engine, CSV_PATH

def get_processed_data():
	with open(CSV_PATH + 'bill_of_materials.csv', 'r') as file:
		reader = csv.DictReader(file)

		rows = []

		for row in reader:
			tube_assembly_id = row['tube_assembly_id']

			for i in range(1,9):

				if row[f'component_id_{i}'] != 'NA':
					component_id = row[f'component_id_{i}']
					quantity 	 = int(row[f'quantity_{i}'])

					rows.append({
						'tube_assembly_id' : tube_assembly_id,
						'component_id' 	   : component_id,
						'quantity'         : quantity
						})
				else:
					pass

	return rows

def add_to_table(session):
	bill_of_materials = get_processed_data()

	for bill_of_material in bill_of_materials:
		row = BillOfMaterials(**bill_of_material)
		session.add(row)

if __name__ == '__main__':
	Session = sessionmaker(bind = oltp_engine)
	session = Session()

	add_to_table(session)
	session.commit()

	print("Success on bill_of_material!")


# Lab 6 Softmax Classifier
from lib.neural_network_one_hot import NeuralNetworkOneHot
from lib.nntype import NNType

class XXX (NeuralNetworkOneHot):
    def init_network(self):
        self.set_placeholder(16, 1)

        self.target_to_one_hot(7)

        # 16 to 7
        logits = self.fully_connected_layer(self.X, 16, 7, 'W', 'b')
        hypothesis = self.softmax(logits)

        self.set_hypothesis(hypothesis)
        self.set_cost_function_with_one_hot(logits, self.get_one_hot()) #not hypothesis, but logits
        self.set_optimizer(NNType.GRADIENT_DESCENT, 0.1)

gildong = XXX()
xdata, ydata = gildong.load_file('data-04-zoo.csv')
gildong.learn(xdata, ydata, 2000, 100)
gildong.print_error()
gildong.evaluate('data-04-zoo.csv')
gildong.show_error()


from typing import Literal

from pydantic import BaseModel

REASON = {
    "BAD_REQUEST": 400,
    "FORBIDDEN": 403,
    "NOT_FOUND": 404,
    "CONFLICT": 409,
    "UNKNOWN": 500,
}

class Error(BaseModel):
    reason: Literal[
        "BAD_REQUEST",
        "NOT_FOUND",
        "CONFLICT",
        "UNKNOWN",
    ]
    message: str

class UserToken(BaseModel):
    id: int
    superuser: bool


"""
加载pb模型
"""
import os

import tensorflow as tf

from data.data_manage import DataManager

pb_file_path = os.getcwd()
sess = tf.Session()
with tf.gfile.FastGFile(pb_file_path + '/pb/model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')  # 导入计算图

    # 需要有一个初始化的过程
    sess.run(tf.global_variables_initializer())
    scope = sess.graph.get_name_scope()
    keys = sess.graph.get_all_collection_keys()
    ops = sess.graph.get_operations()

    input_images = sess.graph.get_tensor_by_name('input_images_batch:0')
    keep_prob22 = sess.graph.get_tensor_by_name('keep_prob22:0')
    keep_prob20 = sess.graph.get_tensor_by_name('keep_prob20:0')
    probability = sess.graph.get_tensor_by_name('probability:0')
    data_manager = DataManager()
    train_batch = data_manager.get_batch_train_data(batch_size=1)
    imgs = data_manager.get_imgs(train_batch)
    full_connection_layer23 = sess.graph.get_tensor_by_name('full_connection_layer23/full_connection_layer23_fc:0')
    ret = sess.run(probability, feed_dict={input_images: imgs, keep_prob22: 0.99, keep_prob20: 0.99})
    full_connection_layer23_values = sess.run(full_connection_layer23,
                                              feed_dict={input_images: imgs, keep_prob22: 0.99, keep_prob20: 0.99})
    print(ret)
    print(full_connection_layer23_values)


from flask import Flask, request, jsonify
from secret_service import SecretService
from secrets_model import SecretModel
from flask_cors import cross_origin

app = Flask(__name__)
secretService = SecretService()

@app.route('/', methods=['GET'])
@cross_origin()
def index():
    return "Welcome to API"

@app.route('/secretos', methods = ['GET'])
@cross_origin()
def getSecrets():
    response =  jsonify(secretService.getSecrets())
    if response is None:
        return jsonify({"response":"Not found", "status":404})
    return response

@app.route('/secretos', methods = ['POST'])
@cross_origin()
def createUser():
    data  = request.get_json()
    secretsModel = SecretModel()
    secretsModel.__dict__ = data
    response = secretService.createSecrets(secretsModel)
    if response is None:
        return jsonify({"response":"Unprocessable Entity", "status":422})
    return jsonify({"responde":"User creado exitosamente!","status":200})

@app.route('/secretos/<id>', methods = ["GET"])
@cross_origin()
def getById(id):
    if not id:
        return None
    response = secretService.getById(id)
    if response is None:
        return jsonify({"response":"Not found", "status":404})
    return jsonify(response)

@app.route('/secretos', methods = ['DELETE'])
@cross_origin()
def deleteAll():
    response = secretService.deleteAll()
    return jsonify({"responde": "Datos eliminado correctamente", "status":200})

if __name__ == '__main__':
    # Threaded option to enable multiple instances for multiple user access support
    app.run(port=5000,debug= True)


import numpy as np
import time
import jax
import jax.numpy as jnp
from jax.experimental import loops
import sys
sys.path.insert(0, "/home/astro/dforero/projects/jax_cosmo") #Use my local jax_cosmo with correlations module
import jax_cosmo as jc
from jax_cosmo.correlations import xicalc_trapz
import numpy as np

import proplot as pplt

from src.mas import cic_mas, cic_mas_vec
from src.correlations import compute_all_correlations, compute_2pt_correlations, bispec, powspec_vec, powspec_vec_fundamental

import MAS_library as MASL
import Pk_library as PKL

box_size = 1000.
redshift = 0.55
b1 = 2.0
seed = 5
create_data = False
if create_data:

    cosmo = cosmology.Planck15
    print(cosmo)
    Plin = cosmology.LinearPower(cosmo, redshift, transfer='EisensteinHu')

    cat = LogNormalCatalog(Plin=Plin, nbar=3.5e-3, BoxSize=box_size, Nmesh=256, bias=b1, seed=seed)
    particles = cat['Position'].compute()
    np.save("data/lognormal.npy", particles)
else:

    particles = np.load("data/lognormal.npy")

particles = jax.device_put(particles)
cosmo = jc.Planck15()

k1, k2 = 0.1, 0.2
theta = jnp.linspace(0, jnp.pi, 20)
s_edges = jnp.arange(1e-3, 200, 5)
s_centers = 0.5 * (s_edges[1:] + s_edges[:-1])
k_edges = jnp.arange(5e-3, 1, 2 * jnp.pi / box_size)

key = jax.random.PRNGKey(5)

n_part = particles.shape[0]
key, subkey = jax.random.split(key)
w = jnp.broadcast_to(jnp.array([1.]), (particles.shape[0],))

shot_noise = box_size**3 / n_part

z = redshift
n_bins = 256
klin = np.logspace(-3, 0, 2048)
plin_i = b1**2 * jc.power.linear_matter_power(cosmo, klin, a = 1. / (1 + z), transfer_fn=jc.transfer.Eisenstein_Hu)
_, xi0, xi2, xi4 = xicalc_trapz(klin, plin_i, 2., s_centers)
theta_theory, Bk_i = PKL.Bispectrum_theory(klin, plin_i, k1, k2)
theta = jnp.linspace(0, jnp.pi, 15)

@jax.jit
def loss(positions):

    #positions += box_size
    #positions %= box_size
    xpos = positions[:,0]
    ypos = positions[:,1]
    zpos = positions[:,2]

    k_ny = jnp.pi * n_bins / box_size

    delta = jnp.zeros((n_bins, n_bins, n_bins))
    delta = cic_mas_vec(delta,
                xpos, ypos, zpos, w,
                n_part,
                0., 0., 0.,
                box_size,
                n_bins,
                True)
    delta /= delta.mean()
    delta -= 1.

    #k, pk, _, s, xi, _, _,_, _, B, Q = compute_all_correlations(delta, box_size, s_edges, k_edges, k1, k2, theta)
    _, _, _, B, Q = bispec(delta, box_size, k1, k2, theta)
    k, pk, _ = powspec_vec(delta, box_size, k_edges)

    plin = b1**2 * jc.power.linear_matter_power(cosmo, k, a = 1. / (1 + z), transfer_fn=jc.transfer.Eisenstein_Hu)  + shot_noise

    #mono_loss = jnp.nanmean((jnp.log10(pk[:,0]) - jnp.log10(plin))**2)
    mono_loss_k = jnp.nanmean((pk[:,0] - plin)**2)
    #mono_loss_s = jnp.nanmean((xi[:,0] - xi0)**2)

    bispec_loss = 1e-9 * jnp.nanmean((B - jnp.interp(theta, theta_theory, Bk_i))**2)

    #return mono_loss #+ 0.01 * quad_loss + 0.01 * hexa_loss
    return mono_loss_k + bispec_loss

from jax.experimental import optimizers

# Setup training loop

x0 = particles
learning_rate = 1.
opt_init, opt_update, get_params = optimizers.adam(learning_rate)
opt_state = opt_init(x0)
@jax.jit
def step(step, opt_state):
    grads = jax.grad(loss)(get_params(opt_state))
    opt_state = opt_update(step, grads, opt_state)
    return opt_state
num_steps = 100
s = time.time()
print("Training...", flush=True)
opt_state = jax.lax.fori_loop(0, num_steps, step, opt_state)
print(f"Training took {time.time() - s} s", flush=True)
new_pos = get_params(opt_state)
new_pos = (new_pos + box_size) % box_size

k_ny = jnp.pi * n_bins / box_size
delta = jnp.zeros((n_bins, n_bins, n_bins))
delta = cic_mas_vec(delta,
            new_pos[:,0], new_pos[:,1], new_pos[:,2], w,
            n_part,
            0., 0., 0.,
            box_size,
            n_bins,
            True)
delta /= delta.mean()
delta -= 1.

k, pk, _, s, xi, _, _,_, _, B, Q = compute_all_correlations(delta, box_size, s_edges, k_edges, k1, k2, theta)
k, pk, modes = powspec_vec_fundamental(delta, box_size)
mask = k < k_ny
k = k[mask]
pk = pk[mask]
s, xi, modes = xi_vec_fundamental(delta, box_size)
mask = s < 200
s = s[mask]
xi = xi[mask]

fig, ax = pplt.subplots(nrows=3, ncols=3, sharex=False, sharey=False)
ax[6].imshow(delta[:,:,:].mean(axis=2), colorbar='r', vmin=-1, vmax = 1.6)
ax[6].format(title='Corrected')
ax[0].plot(k, k*pk[:,0], label='Corrected')
ax[1].plot(k, k*pk[:,1], label='Corrected')
ax[2].plot(k, k*pk[:,2], label='Corrected')

ax[3].plot(s, s**2*xi[:,0], label='Corrected')
ax[4].plot(s, s**2*xi[:,1], label='Corrected')
ax[5].plot(s, s**2*xi[:,2], label='Corrected')

ax[8].plot(theta, B, label='Corrected')

delta = jnp.zeros((n_bins, n_bins, n_bins))
delta = cic_mas_vec(delta,
            particles[:,0], particles[:,1], particles[:,2], w,
            n_part,
            0., 0., 0.,
            box_size,
            n_bins,
            True)
delta /= delta.mean()
delta -= 1.
k, pk, _, s, xi, _, _,_, _, B, Q = compute_all_correlations(delta, box_size, s_edges, k_edges, k1, k2, theta)

ax[0].plot(k, k*pk[:,0], label='Log-normal', ls='--')
ax[1].plot(k, k*pk[:,1], label='Log-normal', ls='--')
ax[2].plot(k, k*pk[:,2], label='Log-normal', ls='--')
ax[3].plot(s, s**2*xi[:,0], label='Log-normal', ls = '--')
ax[4].plot(s, s**2*xi[:,1], label='Log-normal')
ax[5].plot(s, s**2*xi[:,2], label='Log-normal')

ax[8].plot(theta, B, label='Log-normal')

ax[3].plot(s, s**2*xi0, label='linear', ls=':')
ax[4].plot(s, s**2*xi2, label='linear', ls=':')
ax[5].plot(s, s**2*xi4, label='linear', ls=':')

ax[0].plot(klin, klin * plin_i, label='Linear', ls=':')
ax[8].plot(theta_theory, Bk_i, label='Linear', ls=':')

ax[0].format(xscale = 'log', yscale = 'log', xlabel='$k$ [$h$/Mpc]', ylabel='$P(k)$')
ax[0,:].format(xscale = 'log', yscale='linear', xlabel='$k$ [$h$/Mpc]', ylabel='$P(k)$')
ax[7].imshow(delta[:,:,:].mean(axis=2), colorbar='r', vmin=-1, vmax = 1.6)
ax[7].format(title='Log-normal')
ax[0].legend(loc='bottom')
ax[3].legend(loc='bottom')
ax[8].legend(loc='bottom')
ax[3].format(xlabel='$s$', ylabel=r'$s^2\xi$')
fig.savefig("plots/lognormal_bispec.png", dpi=300)


import os
from plum import dispatch
from .internal_representation import File
from .string_operations import (
    format_tex_filename,
    string_find_wrapped_content,
    read_file_lines,
)

def read_extracted_line(line: str, root_directory: str):
    try_input = string_find_wrapped_content(line, "\\input{", "}")
    if try_input is None or line.find("%") != -1:
        return line
    else:
        return read_extracted_file(root_directory, try_input)

@dispatch
def read_extracted_file(root_directory: str, filename: str) -> File:
    lines = read_file_lines(
        os.path.join(root_directory, format_tex_filename(filename)),
        lambda line: read_extracted_line(line, root_directory),
    )
    f = File(root_directory, filename)
    f.lines = lines
    return f


import os
import sys
import argparse
import logging
import torch
import tqdm
from tensorboardX import SummaryWriter
import torchvision
from torchvision import datasets, transforms
import torch.optim as optim
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import torchvision.models.resnet as resnet
import numpy as np
from easydict import EasyDict

sys.path.append(os.getcwd())
from utils.utils import BestRecoder, Timer

parser = argparse.ArgumentParser(description="MNIST Resnet pytorch")
parser.add_argument("--batch_size", type=int, default=64, metavar='N', help="batch size (default 64)")
parser.add_argument("--learning_rate", type=float, default=1e-2, metavar="LR", help="learning rate for training, default 10^-2")
parser.add_argument("--no_cuda",action="store_true",default=False, help="disable the cuda")
parser.add_argument("--epoches", type=int, default=30, help="set the training epoches, default 10")
parser.add_argument("--root", type=str, default="/home/lijiguo/pytorch_examples",
                    help="root path for 'pytorch_examples', default '/home/jiguo/pytorch_examples'")
parser.add_argument("--order", type=int, default=1, help="order for Tylor net")
parser.add_argument("--lr_decay", type=float, default=0.1, help="decay each epoch")
parser.add_argument("--weight_decay", type=float, default=4e-5, help="weight decay for optimizer")
parser.add_argument("--block_type", type=str, default="res", help="block type")
parser.add_argument("--log_file", type=str, default="mnist_resnet.log", help="log file")
parser.add_argument("--dataset", type=str, default="mnist", help="dataset, mnist or cifar10")

# args
args = parser.parse_args()

args.cuda = not args.no_cuda and torch.cuda.is_available()

args.root = os.getcwd()

runtime_cfg = EasyDict()

runtime_cfg.model_dir = os.path.split(os.path.dirname(__file__))[-1]  # "./mnist_resnet"
runtime_cfg.log_dir = os.path.join(args.root, runtime_cfg.model_dir, "./log")
logging.basicConfig(format='%(message)s', level=logging.INFO, handlers=[
    logging.FileHandler(os.path.join(runtime_cfg.log_dir, args.log_file)),
    logging.StreamHandler(sys.stdout)
] )
runtime_cfg.logger = logging.getLogger()

if not os.path.exists(runtime_cfg.log_dir):
    os.mkdir(runtime_cfg.log_dir)

runtime_cfg.writer = SummaryWriter(log_dir=runtime_cfg.log_dir)

# 3*3 convolution
def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3,
                    stride=stride, padding=1, bias=False)

# Residual block
class ResidualBlock(nn.Module):
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

class PolyBlock(nn.Module):
    """
    x_{n+1} = x_n + f(x_{n+1}), one conv for f
    """
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(PolyBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = self.conv1(x)
        out = self.bn1(out)
        # out = self.relu(out)

        out1 = self.relu(residual + out)

        out1 = self.conv2(out1)
        out1 = self.bn2(out1)
        out = self.relu(out1 + residual)
        return out

class PolyBlockV1(nn.Module):  # implicit Euler
    """
    x_{n+1} = x_n + f(x_{n+1}), two conv for f
    """
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(PolyBlockV1, self).__init__()
        self.res_conv1 = nn.Sequential(
            conv3x3(in_channels, out_channels, stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_conv2 = nn.Sequential(
            conv3x3(out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = self.res_conv1(x)
        # out = self.relu(out)

        out1 = self.relu(residual + out)

        out1 = self.res_conv2(out1)
        out = self.relu(out1 + residual)
        return out

class PolyBlockV2(nn.Module):  # implicit Euler
    """
    x_{n+1} = x_n + (f(x_n)+f(x_{n+1}))/2
    """
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(PolyBlockV2, self).__init__()
        self.res_conv1 = nn.Sequential(
            conv3x3(in_channels, out_channels, stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_conv2 = nn.Sequential(
            conv3x3(out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = self.res_conv1(x)
        # out = self.relu(out)

        out1 = self.relu(residual + out)

        out1 = self.res_conv2(out1)
        out_final = self.relu((out1+out)/2 + residual)
        return out_final

class PolyBlockV3(nn.Module):  # implicit Euler
    """
    x_{n+1} = x_n + (f(x_n)+f(x_{n+1}))/2
    """
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(PolyBlockV3, self).__init__()
        self.res_conv1 = nn.Sequential(
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_conv2 = nn.Sequential(
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = self.res_conv1(residual)
        # out = self.relu(out)

        out1 = self.relu(residual + out)

        out1 = self.res_conv2(out1)
        out_final = self.relu((out1+out)/2 + residual)
        return out_final

class RKBlock(nn.Module):
    """
    Runge-Kutta
    """
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(RKBlock, self).__init__()
        self.res_conv = nn.Sequential(
            conv3x3(in_channels, out_channels, stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_conv_h_2 = nn.Sequential(
            conv3x3(out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_conv_h = nn.Sequential(
            conv3x3(out_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        K1 = self.res_conv(x)
        K2 = self.res_conv_h_2(self.relu(residual+K1/2))
        K3 = self.res_conv_h_2(self.relu(residual+K2/2))
        K4 = self.res_conv_h(self.relu(residual+K3))

        out_final = self.relu(residual+K1/6 + K2/3 + K3/3 + K4/6)

        return out_final

class TylorBlockV1(nn.Module):  # Tylor
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(TylorBlockV1, self).__init__()
        self.res_conv1 = nn.Sequential(
            conv3x3(in_channels, out_channels, stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            conv3x3(out_channels, out_channels),
            nn.BatchNorm2d(out_channels),
        )
        self.res_convs = []
        self.order = args.order
        for idx in range(self.order-1):
            block = nn.Sequential(
                    conv3x3(out_channels, out_channels, 1),
                    nn.BatchNorm2d(out_channels),
                    nn.ReLU(inplace=True),
                    conv3x3(out_channels, out_channels),
                    nn.BatchNorm2d(out_channels),
                    )
            if args.cuda:
                block = block.cuda()
            self.res_convs.append(
                block
            )
        self.relu = nn.ReLU(inplace=True)

        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = self.res_conv1(x)
        # out = self.relu(out)

        # tylors = [residual, out]
        base = 1
        out_final = residual + out
        for idx in range(self.order-1):
            base *= (idx+1)
            out = self.res_convs[idx](self.relu(out))
            out_final += (out/base)
        out_final = self.relu(out_final)
        return out_final

class DenseBlock(nn.Module):  # Dense
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(DenseBlock, self).__init__()
        self.res_convs = []
        self.order = args.order
        for idx in range(self.order):
            block = nn.Sequential(
                    conv3x3(out_channels, out_channels, 1),
                    nn.BatchNorm2d(out_channels),
                    nn.ReLU(inplace=True),
                    conv3x3(out_channels, out_channels),
                    nn.BatchNorm2d(out_channels),
                    )
            if args.cuda:
                block = block.cuda()
            self.res_convs.append(
                block
            )
        self.relu = nn.ReLU(inplace=True)

        self.downsample = downsample

    def forward(self, x):
        residual = x
        if self.downsample:
            residual = self.downsample(x)

        out = residual
        # out = self.relu(out)

        # tylors = [residual, out]
        out_final = residual
        for idx in range(self.order):
            out = self.res_convs[idx](self.relu(out))
            out_final += (out)
        out_final = self.relu(out_final)
        return out_final

# ResNet
class ResNet_Mnist(nn.Module):
    def __init__(self, block, layers, num_classes=10, ndim=16):
        super(ResNet_Mnist, self).__init__()
        self.in_channels = ndim
        self.conv = conv3x3(args.input_channel, ndim)
        self.bn = nn.BatchNorm2d(ndim)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self.make_layer(block, ndim, layers[0])
        self.layer2 = self.make_layer(block, ndim*2, layers[1], 2)
        self.layer3 = self.make_layer(block, ndim*4, layers[2], 2)
        self.layer4 = self.make_layer(block, ndim*8, layers[3], 2)
        self.avg_pool = nn.AvgPool2d(3)
        self.fc = nn.Linear(ndim*8, num_classes)

    def make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels):
            downsample = nn.Sequential(
                conv3x3(self.in_channels, out_channels, stride=stride),
                nn.BatchNorm2d(out_channels))
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

if args.dataset == "mnist":
    train_dataset = torchvision.datasets.MNIST(root=os.path.join(args.root, "./data/mnist"),train=True, download=True,
                transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))
                        ]))
    test_dataset = torchvision.datasets.MNIST(root=os.path.join(args.root, "./data/mnist"),train=False,
                transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))
                        ]))
    args.input_channel = 1
    args.num_classes = 10
    ResNet = ResNet_Mnist
elif args.dataset == "cifar10":
    train_dataset = torchvision.datasets.CIFAR10(root=os.path.join(args.root, "./data/cifar10"),train=True, download=True,
                transform=transforms.Compose([
                        transforms.RandomCrop(32, padding=4),
                        transforms.RandomHorizontalFlip(),
                        transforms.ToTensor(),
                        transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010))
                        ]))
    test_dataset = torchvision.datasets.CIFAR10(root=os.path.join(args.root, "./data/cifar10"),train=False,
                transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010))
                        ]))
    args.input_channel = 3
    args.num_classes = 10
    ResNet = ResNet_Mnist
else:
    raise NotImplementedError

param = {"num_workers": 8, "pin_memory": True} if args.cuda else {"num_workers": 8}
train_loader=torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **param)
test_loader=torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, **param)

def l1_loss(params):
    param_sum = torch.zeros(1)
    param_cnt = 0
    for param in params:
        param_sum += torch.pow(param,2).mean()*0.5
        param_cnt += 1
    return param_sum/param_cnt

def my_loss(pred, target):
    loss = -pred + target*torch.log(pred) - (1-target)*torch.log(1-pred)
    return loss

blocks = {
    "res": ResidualBlock,
    "poly": PolyBlock,
    "poly1": PolyBlockV1,
    "RK": RKBlock,
    "poly2": PolyBlockV2,
    "poly3": PolyBlockV3,
    "tylor": TylorBlockV1
}

# Create ResNet
net_args = {
    "block": blocks[args.block_type], #PolyBlockV1, # ResidualBlock,
    "layers": [2, 2, 2, 2],
    "num_classes": args.num_classes,
    "ndim": 16 if args.dataset=='mnist' else 64
}
model = ResNet(**net_args)
if args.cuda:
    model = model.cuda()

optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, momentum=0.9)
loss_func = torch.nn.CrossEntropyLoss()

runtime_cfg.lr_decay_epoches = [int(args.epoches//4), int(args.epoches//2), int(args.epoches*3//4)]
runtime_cfg.current_lr = args.learning_rate

def train(epoch, runtime_cfg):
    """
    param:
    return:
    """
    model.train()
    train_loss = 0
    loader_bar = tqdm.tqdm(train_loader)
    # lr decay
    if epoch in runtime_cfg.lr_decay_epoches:
        runtime_cfg.current_lr *= args.lr_decay
        for param_group in optimizer.param_groups:
            param_group['lr'] = runtime_cfg.current_lr
        runtime_cfg.logger.info("decay lr to {}".format(runtime_cfg.current_lr))
    for batch_idx, (data, label) in enumerate(loader_bar):
        data, label = data.requires_grad_(), label.requires_grad_(False)
        if args.cuda:
            data, label = data.cuda(), label.cuda()

        torch.unsqueeze(data, 2)
        pred = model.forward(data)
        loss = loss_func(pred, label) #+ 100 * l1_loss(model.linear1.parameters())

        model.zero_grad()
        loss.backward()
        optimizer.step()
        loader_bar.set_description("[{:5s}] Epoch:{:5d}, loss:{:10.5f}".format("train", epoch, loss.cpu().item()))
        train_loss += loss.cpu().item()
    loader_bar.close()
    train_loss = train_loss/len(train_loader)
    runtime_cfg.logger.info("[{:5s}] Epoch:{:5d}, loss:{:10.5f}".format("train", epoch, train_loss))

def test(epoch, runtime_cfg):
    model.eval()
    test_loss = 0
    accuracy = 0
    loader_bar = tqdm.tqdm(test_loader)
    for batch_idx, (data, label) in enumerate(loader_bar):
        data, label = data.requires_grad_(False), label.requires_grad_(False)
        if args.cuda:
            data, label = data.cuda(), label.cuda()

        torch.unsqueeze(data, 2)
        pred = model.forward(data)
        loss = loss_func(pred, label) #+ l1_loss(model.linear1.parameters())
        test_loss += loss.cpu().item()
        pred_label = pred.data.max(1,keepdim=True)[1]
        accuracy += pred_label.eq(label.data.view_as(pred_label)).long().cpu().sum()
        loader_bar.set_description("[{:5s}] Epoch:{:5d}, loss:{:10.5f}".format("test", epoch, loss.cpu().item()))

    accuracy = float(accuracy)/len(test_loader.dataset)
    test_loss = test_loss/len(test_loader)

    loader_bar.close()
    runtime_cfg.logger.info("[{:5s}] Epoch:{:5d}, loss:{:10.5f}, accuracy:{:10.5f}".format("test", epoch, test_loss, accuracy))

    model_path = os.path.join(args.root, runtime_cfg.model_dir, "./model")
    if not os.path.exists(model_path):
        logging.warning("path ({}) does not exist, create it".format(model_path))
        os.system("mkdir {}".format(model_path))

    torch.save(model.state_dict(),os.path.join(model_path, "./model_{}.pytorch".format(epoch)))
    return accuracy

if __name__ == "__main__":
    runtime_cfg.logger.info(args)
    best_recoder = BestRecoder(init=0)
    for epoch in range(0, args.epoches):
        train(epoch, runtime_cfg)
        accu = test(epoch, runtime_cfg)
        best_recoder.update(accu)

    runtime_cfg.logger.info("best accu: {:10.5f}".format(best_recoder.best))


import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error
from sklearn.impute import SimpleImputer

def score_dataset(X_train, X_valid, y_train, y_valid):
    model = RandomForestRegressor(n_estimators=50, random_state=0)
    model.fit(X_train, y_train)
    preds = model.predict(X_valid)
    return mean_absolute_error(y_valid, preds)

def main():
    X = pd.read_csv('train.csv', index_col='PassengerId')
    X_test = pd.read_csv('test.csv', index_col='PassengerId')
    y = X.Survived
    X.drop(['Survived'], axis=1, inplace=True)

    # # Split training and test data.
    X_train, X_valid, y_train, y_valid = train_test_split(X, y, random_state=0)

    # Which columns have missing values?
    cols_with_missing = [col for col in X.columns
                        if X[col].isnull().any()]
    print('Columns with missing values:')
    print(cols_with_missing)

    # Numeric columns are:
    # - Pclass (ticket class)
    # - Age
    # - SibSp (# of siblings and spouses aboard)
    # - Parch (# of parents and children aboard)
    # - Fare
    numeric_cols = [col for col in X.columns
                   if X[col].dtype in ['float64', 'int64']]
    print('Numeric columns:')
    print(numeric_cols)

    # Categorical columns
    categorical_cols = [col for col in X.columns
                       if X[col].dtype in ['object']]
    print('Categorical columns:')
    print(categorical_cols)

    # Use only numeric columns to start
    numeric_X_train = X_train[numeric_cols].copy()
    numeric_X_valid = X_valid[numeric_cols].copy()

    print(numeric_X_train)
    print(type(numeric_X_train))

    # Delete Age column.
    reduced_X_train = numeric_X_train.drop(['Age'], axis=1)
    reduced_X_valid = numeric_X_valid.drop(['Age'], axis=1)

    print(reduced_X_train)
    print(type(reduced_X_train))

    print('MAE (delete Age column):')
    print(score_dataset(reduced_X_train, reduced_X_valid, y_train, y_valid))

    # Replace missing age values with mean age.
    my_imputer = SimpleImputer(strategy='mean')
    imputed_X_train = pd.DataFrame(my_imputer.fit_transform(numeric_X_train))
    imputed_X_valid = pd.DataFrame(my_imputer.transform(numeric_X_valid))
    imputed_X_train.columns = numeric_X_train.columns
    imputed_X_valid.columns = numeric_X_valid.columns

    print('MAE (impute Age column):')
    print(score_dataset(imputed_X_train, imputed_X_valid, y_train, y_valid))

if __name__ == '__main__':
    main()


import argparse
import os
import sys

import reklib

LEVELS = (
    ("pusy", "use lzma"),
    ("chiken", "lossless, keep useless data"),
    ("norm", "lossless"),
    ("oke", "lossless, discard sounds"),
    ("very", "lossy"),
    ("max", "lossy, discard sounds")
)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("-f", "--force", help="overwrite files silently",
                        action="store_true")
    parser.add_argument("infile", nargs="?")
    parser.add_argument("outfile", nargs="?")

    group = parser.add_argument_group("compression levels")
    group = group.add_mutually_exclusive_group()
    group.add_argument("-m", "--moar", help="compress harder",
                       dest="level", action="count", default=0)
    for i, (name, desc) in enumerate(LEVELS):
        group.add_argument(f"-{i}", f"--{name}", help=desc,
                           dest="level", action="store_const", const=i)

    args = parser.parse_args()

    try:
        if args.infile:
            if not args.outfile:
                args.outfile = args.infile + ".rek"
            if os.path.exists(args.outfile) and not args.force:
                print(f"Cowardly refusing to overwrite {args.outfile}", file=sys.stderr)
                exit(1)

            reklib.compress(open(args.infile, "rb"), open(args.outfile, "wb"), args.level)

        elif not sys.stdin.isatty():
            reklib.compress(sys.stdin.buffer, sys.stdout.buffer, args.level)

        else:
            parser.print_help()

    except NotImplementedError:
        print("orka", file=sys.stderr)
        exit(1)


import bpy
from gorgious_utilities.collection.tool import (
    get_family_down,
)
from gorgious_utilities.custom_property.helper import remove_trailing_numbers

class GU_OT_collection_rename_objects(bpy.types.Operator):
    """Rename each object in the collection"""

    bl_idname = "collection.rename_objects"
    bl_label = "Rename each object in the collection"
    bl_options = {"UNDO", "REGISTER"}

    replace_from: bpy.props.StringProperty(name="Replace")
    replace_to: bpy.props.StringProperty(name="With")
    as_collection: bpy.props.BoolProperty(name="Rename As Collection", default=True)
    recursive: bpy.props.BoolProperty(
        name="Recursive", description="Replace name in collection's children hierarchy", default=True
    )
    remove_trailing_numbers: bpy.props.BoolProperty(default=False, options={"SKIP_SAVE"})

    def draw(self, context):
        layout = self.layout
        layout.prop(self, "recursive")
        layout.prop(self, "as_collection")
        col = layout.column()
        col.prop(self, "replace_from")
        col.prop(self, "replace_to")
        col.enabled = not self.as_collection

    def invoke(self, context, event):
        if self.remove_trailing_numbers:
            return self.execute(context)
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        cols = get_family_down(context.collection) if self.recursive else (context.collection,)
        for col in cols:
            if self.remove_trailing_numbers:
                remove_trailing_numbers(col, "name")
            for obj in col.objects:
                if self.remove_trailing_numbers:
                    remove_trailing_numbers(obj, "name")
                else:
                    obj.name = col.name if self.as_collection else obj.name.replace(self.replace_from, self.replace_to)
                if hasattr(obj, "data") and obj.data is not None:
                    obj.data.name = obj.name

        return {"FINISHED"}


"""This script will run the avalanche model and save the number of defaults in a CSV with columns for Ratio Defaults, Cascade Defaults
and rows representing the step. From that, we can plot the timeline of defaults."""

import json
from tqdm import tqdm
import numpy as np
import cvxpy as cvx
import time
from contagion import binarize_probabilities, distribute_liabilities, make_connections, DeterministicRatioNetwork, TestNetwork

# adjust numberOfRuns to change number of times entire model is run
# default = 1
numberOfRuns = 100
# adjust steps to change the number of steps per run of the model
# default = 1000000
steps = 1000000

# change distribution:
# options:
# normal     - Change distribution variable to 'normal'
# poisson    - Change distribution variable to 'poisson'
# lognormal  - Change distribution variable to 'lognormal'
# gamma      - Change distribution variable to 'gamma'

#### Distribution variables:
# used for all
cashDistribution = 'poisson'
leverageDistribution = 'poisson'
size = 100;  # ## DO NOT CHANGE!

# # normal ##
  # cash #
normalCashLocation = 10000;  # default = 10000
normalCashScale = 10000;  # default = 10000
  # leverage #
normalLeverageLocation = 10;  # default = 10
normalLeverageScale = 2;  # default = 2

# # poisson ##
  # cash #
poissonCashLambda = 3;  # default = 3
  # leverage #
poissonLeverageLambda = 3;  # default = 3

# # lognormal ##
  # cash #
lognormalCashMean = 0;  # default = 0
lognormalCashSigma = 1;  # default = 1
  # leverage #
lognormalLeverageMean = 0;  # default = 0
lognormalLeverageSigma = 1;  # default = 1

# # gamma ##
  # cash #
gammaCashShape = 2;  # default = 2
gammaCashScale = 1;  # default = 1
  # leverage #
gammaLeverageShape = 2;  # default = 2
gammaLeverageScale = 1;  # default = 1

def runModel(cashDistribution, leverageDistribution):
    cash_vector = generateCashVector(cashDistribution)
    cash_vector[cash_vector <= 0] = 1 * 10 ** -10
    # cash_vector[cash_vector > 5000] = 6500
    cash_to_connectivity = lambda x: safe_ln(x)
    connectivity_vector = cash_to_connectivity(cash_vector)

    # Make the adjacency matrix
    mat = make_connections(connectivity_vector)
    mat = binarize_probabilities(mat)

    # Distribute liabilities
    leverage_ratios = generateLeverageRatios(leverageDistribution)
    leverage_ratios[leverage_ratios < 5] = 5

    liabilities = np.multiply(cash_vector, leverage_ratios)
    mat = distribute_liabilities(mat, liabilities)
    for i, cash in enumerate(cash_vector):
        mat[i, i] = cash
    defaults = np.zeros((steps, 2))
    for z in tqdm(range(steps)):
        model = TestNetwork(100, mat)
        model.reset_net()

        results = model.step()
        defaults[z, 0] = results['ratio_defaults']
        defaults[z, 1] = results['cascade_defaults']

    """ the below code saves the results in a .csv file with the date and time added to
    the name so as to prevent overwriting. The format parameter can be used
    to change the variable type that is stored in the .csv file, with the result being
    smaller file sizes when more data efficient variable types are utilized."""

    # use fmt='%1c' to save data as characters
    # use fmt='%1d' to save data as signed decimal integers
    # use fmt='%1e' to save data in scientific notation
    # use fmt='%1f' to save data as decimal floating points
    # use fmt='%1g' to save data as the most efficient of e or f
    # use fmt='%1o' to save data as signed octals
    # use fmt='%1s' to save data as string characters
    # use fmt='%1u' to save data as unsigned decimal integers
    # use fmt='%1x' to save data as unsigned hexadecimal integer

    np.savetxt('defaults_' + str(time.strftime("%d_%m_%y_%H%M%S")) + '.csv', defaults, delimiter=',', fmt='%1f')

    """ Cursory analysis reveals that with the standard settings and no fmt parameter used
    the data files are approximately 50MB in size. Manually formatting in excel can reduce
    the file size to 5MB.

    Using signed decimal integers the file size is approximately 4MB
    Using decimal floating points the file size is approximately 18MB"""

# The below function generates the chosen cash distribution
def generateCashVector(distribution):
    if distribution == 'normal':
        cash_vector = np.random.normal(normalCashLocation, normalCashScale, size)
    elif distribution == 'poisson':
        cash_vector = np.random.poisson(poissonCashLambda, size)
    elif distribution == 'lognormal':
        cash_vector = np.random.lognormal(lognormalCashMean, lognormalCashSigma, size)
    elif distribution == 'gamma':
        cash_vector = np.random.gamma(gammaCashShape, gammaCashScale, size)
    else: cash_vector = np.random.normal(normalCashLocation, normalCashScale, size)
    return cash_vector

# The below function generates a string to reflect the chosen cash distribution
def generateCashString(distribution):
    cashString = 'asdf'
    if distribution == 'normal':
        cashString = 'NormalCash_location' + str(normalCashLocation) + 'Scale' + str(normalCashScale) + '_'
    elif distribution == 'poisson':
        cashString = 'PoissonCash_Lambda' + str(poissonCashLambda) + '_'
    elif distribution == 'lognormal':
        cashString = 'LognormalCash_Mean' + str(lognormalCashMean) + 'Sigma' + str(lognormalCashSigma) + '_'
    elif distribution == 'gamma':
        cashString = 'GammaCash_Shape' + str(gammaCashShape) + 'Scale' + str(gammaCashScale) + '_'
    else: cashString = 'NormalCash_location' + str(normalCashLocation) + 'Scale' + str(normalCashScale) + '_'
    return cashString

# The below function generates the chosen leverage distribution
def generateLeverageRatios(distribution):
    if distribution == 'normal':
        leverage_ratios = np.random.normal(normalLeverageLocation, normalLeverageScale, size)
    elif distribution == 'poisson':
        leverage_ratios = np.random.poisson(poissonLeverageLambda, size)
    elif distribution == 'lognormal':
        leverage_ratios = np.random.lognormal(lognormalLeverageMean, lognormalLeverageSigma, size)
    elif distribution == 'gamma':
        leverage_ratios = np.random.gamma(gammaLeverageShape, gammaLeverageScale, size)
    else: leverage_ratios = np.random.normal(normalLeverageLocation, normalLeverageScale, size)
    return leverage_ratios

# The below function generates a string to reflect the chosen leverage distribution
def generateLeverageString(distribution):
    leverageString = 'asdf'
    if distribution == 'normal':
        leverageString = 'NormalLeverage_location' + str(normalLeverageLocation) + 'Scale' + str(normalLeverageScale) + '_'
    elif distribution == 'poisson':
        leverageString = 'PoissonLeverage_Lambda' + str(poissonLeverageLambda) + '_'
    elif distribution == 'lognormal':
        leverageString = 'LognormalLeverage_Mean' + str(lognormalLeverageMean) + 'Sigma' + str(lognormalLeverageSigma) + '_'
    elif distribution == 'gamma':
        leverageString = 'GammaLeverage_Shape' + str(gammaLeverageShape) + 'Scale' + str(gammaLeverageScale) + '_'
    else: leverageString = 'NormalLeverage_location' + str(normalLeverageLocation) + 'Scale' + str(normalLeverageScale) + '_'
    return leverageString

# The below function checks for and corrects potential division by zero
def safe_ln(x, minval=0.000000000001):
    return np.log(x.clip(min=minval)).astype(int)

''' Create strings to use in saving results so there is record of distributions
    used and their parameters: '''
cashString = generateCashString(cashDistribution)
leverageString = generateLeverageString(leverageDistribution)

# the below loop runs the program for the desired number of iterations
for j in range(numberOfRuns):
    runModel(cashDistribution, leverageDistribution)


from cloud_blobstore import BlobStore
from cloud_blobstore import BlobNotFoundError
from enum import Enum, auto

class ObjectTest(Enum):
    EXACT = auto()
    PREFIX = auto()

def test_object_exists(blobstore: BlobStore, bucket: str, match: str, test_type: ObjectTest = ObjectTest.EXACT) -> bool:
    """
    Test if an object exists in the BlobStore
    :param blobstore: the blobstore to check
    :param bucket: the bucket to check in the blobstore
    :param match: the string to match against; this is _not_ a regex pattern, strings must match exactly
    :param test_type: the type of test to conduct, prefix matches test if the object name starts with the match string,
        exact matches must match the full string
    :return: test bool
    """
    if test_type == ObjectTest.PREFIX:
        try:
            blobstore.list(bucket, prefix=match).__iter__().__next__()
        except StopIteration:
            return False
        else:
            return True
    elif test_type == ObjectTest.EXACT:
        try:
            blobstore.get_user_metadata(bucket, match)
            return True
        except BlobNotFoundError:
            return False
    else:
        raise ValueError(f"Not a valid storage object test type: " + test_type.name)