import arcpy
import os
import random
import string
import tempfile
import unittest
import arcgisscripting
from os import path

def get_temp_name():
    """
    Creates a random string that is a safe table or feature class name
    :return:
    """
    return ''.join([random.choice(string.ascii_uppercase) for _ in range(6)])

class FeatureClassAssertMixin(object):
    """
    Mixin for supporting the compare or assertion of two feature classes.
    """

    def assertFeatureClassEqual(self,
                                first,
                                second,
                                sort_field,
                                message=None,
                                compare_type=None,
                                ignore_options=None,
                                xy_tolerance=None,
                                m_tolerance=None,
                                z_tolerance=None,
                                attribute_tolerances=None,
                                omit_field=None):
        """
        Compares the second feature class to the first, and reports any issues.  Detailed issues are printed to the
        console if found.
        :param first:
        :param second:
        :param sort_field:
        :param message:
        :param compare_type:
        :param ignore_options:
        :param xy_tolerance:
        :param m_tolerance:
        :param z_tolerance:
        :param attribute_tolerances:
        :param omit_field:
        :return:
        """

        # Make a place to store the compare file
        compare_file = tempfile.mkstemp(".txt")
        os.close(compare_file[0])
        result = arcpy.FeatureCompare_management(first,
                                                 second,
                                                 sort_field,
                                                 compare_type,
                                                 ignore_options,
                                                 xy_tolerance,
                                                 m_tolerance,
                                                 z_tolerance,
                                                 attribute_tolerances,
                                                 omit_field,
                                                 continue_compare=True,
                                                 out_compare_file=compare_file[1])
        if 'true' == result.getOutput(1):
            # delete the compare file
            os.remove(compare_file[1])
            return

        # read the compare file and print it out
        print(compare_file[1])
        with open(compare_file[1]) as f:
            [print(l.rstrip()) for l in f.readlines()]
        os.remove(compare_file[1])

        # set the assertion message.
        msg = message if message is not None else "Feature class {} is not equal to {}".format(second, first)
        raise AssertionError(msg)

class SimpleToolboxTests(unittest.TestCase, FeatureClassAssertMixin):

    def setUp(self):
        arcpy.env.overwriteOutput = True
        self.parent_directory = path.dirname(path.expandvars(__file__))
        self.data_directory = path.join(self.parent_directory, "data")
        arcpy.AddToolbox(path.join(path.dirname(self.parent_directory), "SimpleToolbox.pyt"))

    def test_single_polygon_red(self):
        """
        A single polygon with three red points.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "single_polygon_red.gdb", "points")
        polygons = path.join(self.data_directory, "single_polygon_red.gdb", "polygons")
        answer = path.join(self.data_directory, "single_polygon_red.gdb", "answer")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run
        arcpy.SimpleTool_simple(points, polygons, output)
        # assert
        self.assertFeatureClassEqual(answer, output, "OBJECTID")
        # cleanup
        arcpy.Delete_management(output)

    def test_single_polygon_blue(self):
        """
        A single polygon with three blue points.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "single_polygon_blue.gdb", "points")
        polygons = path.join(self.data_directory, "single_polygon_blue.gdb", "polygons")
        answer = path.join(self.data_directory, "single_polygon_blue.gdb", "answer")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run
        arcpy.SimpleTool_simple(points, polygons, output)
        # assert
        self.assertFeatureClassEqual(answer, output, "OBJECTID")
        # cleanup
        arcpy.Delete_management(output)

    def test_single_polygon_no_points(self):
        """
        A single polygon with no points in the points layer (empty, no records)
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "single_polygon_no_points.gdb", "points")
        polygons = path.join(self.data_directory, "single_polygon_no_points.gdb", "polygons")
        answer = path.join(self.data_directory, "single_polygon_no_points.gdb", "answer")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run
        arcpy.SimpleTool_simple(points, polygons, output)
        # assert
        self.assertFeatureClassEqual(answer, output, "OBJECTID")
        # cleanup
        arcpy.Delete_management(output)

    def test_single_polygon_no_points_by_filter(self):
        """
        A single polygon with no points in the points layer that match either BLUE or RED, but there are ORANGE ones.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "single_polygon_no_points_by_filter.gdb", "points")
        polygons = path.join(self.data_directory, "single_polygon_no_points_by_filter.gdb", "polygons")
        answer = path.join(self.data_directory, "single_polygon_no_points_by_filter.gdb", "answer")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run
        arcpy.SimpleTool_simple(points, polygons, output)
        # assert
        self.assertFeatureClassEqual(answer, output, "OBJECTID")
        # cleanup
        arcpy.Delete_management(output)

    def test_random_data_01(self):
        """
        A 10x10 grid with 100 random points.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "random_data_01.gdb", "points")
        polygons = path.join(self.data_directory, "random_data_01.gdb", "polygons")
        answer = path.join(self.data_directory, "random_data_01.gdb", "answer")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run
        arcpy.SimpleTool_simple(points, polygons, output)
        # assert
        self.assertFeatureClassEqual(answer, output, "OBJECTID")
        # cleanup
        arcpy.Delete_management(output)

    def test_invalid_point_data(self):
        """
        Passes a polygon in to the point layer.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "random_data_01.gdb", "polygons")
        polygons = path.join(self.data_directory, "random_data_01.gdb", "polygons")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run with assert exception raised
        with self.assertRaisesRegex(arcgisscripting.ExecuteError, 'Invalid geometry type'):
            arcpy.SimpleTool_simple(points, polygons, output)

    def test_invalid_polygon_data(self):
        """
        Passes a point in to the polygon layer.
        :return:
        """
        # Setup
        points = path.join(self.data_directory, "random_data_01.gdb", "points")
        polygons = path.join(self.data_directory, "random_data_01.gdb", "points")
        output = path.join("%ScratchGDB%", get_temp_name())
        # Run with assert exception raised
        with self.assertRaisesRegex(arcgisscripting.ExecuteError, 'Invalid geometry type'):
            arcpy.SimpleTool_simple(points, polygons, output)

    # def test_invalid_point_layer_missing_type_field(self):
    #     """
    #     Passes a point layer that is missing the type field.
    #     :return:
    #     """
    #     # Setup
    #     points = path.join(self.data_directory, "point_layer_no_type_field.gdb", "points")
    #     polygons = path.join(self.data_directory, "point_layer_no_type_field.gdb", "polygons")
    #     output = path.join("%ScratchGDB%", get_temp_name())
    #     # Run with assert exception raised
    #     with self.assertRaisesRegex(arcgisscripting.ExecuteError, 'does not contain a string field called TYPE'):
    #         arcpy.SimpleTool_simple(points, polygons, output)

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


#listings app-urls.py
from django.urls import path

from . import views

urlpatterns = [
    #the name is the page link we use {% urls '' %}
    path('', views.index, name='listings'),
    path('<int:listing_id>', views.listing, name='listing'),
    path('search', views.search, name='search'),
]


from apiclient.discovery import build
from httplib2 import Http
from oauth2client import file, client, tools
import os

"""
Convenience class wrapper for Google Drive OAuth.

This requires that the user download credentials.json
from Google Cloud Console (API section).

The user should set client_secret_file to that JSON file,
probably named client_secret.json, that contains the
application's public and private authentication token.
These tokens are APPLICATION specific.

The authentication process uses these credentials to ask for
an OAuth token on behalf of the user. This step requires the
user to log in, and when they do it creates credentials.json.
These contain the credentials to perform actions as the user
that just logged in/granted permission to your app.
"""

class GDrive(object):
    def __init__(self, gdrive_token_path, config):
        """
        Set up the Google Drive API instance.
        Create the API instance here, hand it over in get_service().
        """
        # Setup the Drive v3 API
        self.SCOPES = 'https://www.googleapis.com/auth/drive.metadata.readonly'
        self.store = file.Storage(gdrive_token_path)

    def get_service(self):
        """
        Return an instance of the Google Drive API service.
        """
        creds = self.store.get()
        if not creds or creds.invalid:
            raise Exception("Error: invalid or missing Google Drive API credentials")

        service = build('drive', 'v3', http=creds.authorize(Http()))
        return service

if __name__=="__main__":
    g = GDrive()
    s = g.get_service()


# 将函数存储在模块中
"""函数的优点之一是，使用他们可将代码块与主程序分离。通过给函数指定描述性名称，
可让主程序容易理解得多。你还可以更进一步，将函数存储在被称为模块的独立文件中，
再将模块导入到主程序中。import语句允许在当前运行的程序文件中使用模块中的代码。
通过将函数存储在堵路的文件中，可隐藏程序代码的细节，将重点放在程序的高层逻辑上。
这还能让你在众多不同的程序中重用函数。将函数存储在独立文件中后，可与其他程序员共享
这些文件而不是整个程序。知道如何导入函数还能让你使用其他程序员编写的函数库。"""

# 导入整个模块
"""要让函数是可导入的，得先创建模块。模块是扩展名为.py的文件，
包含要导入到程序中的代码。
"""

def make_pizza(size, *toppings):
    """概述要制作的比萨"""
    print("\nMaking a " + str(size) +
          "-inch pizza with the following toppings:")
    for topping in toppings:
        print("- " + topping)


from typing import Tuple

class Paddle(object):
    """The Paddle object represents the paddle used in pong game.

    :param x: The current x position of the paddle on the game board.
    :param y: The current y position of the paddle on the game board.
    :param width: The width of the paddle.
    :param height: The height of the paddle.
    """
    VELOCITY = 4

    def __init__(self, x: int, y: int, width: int, height: int) -> None:
        self.original_x = x
        self.original_y = y
        self.x = x
        self.y = y
        self.width = width
        self.height = height

    def move(self, right: bool=True) -> None:
        if right:
            self.x += self.VELOCITY
        else:
            self.x -= self.VELOCITY

    def reset(self) ->None:
        self.x = self.original_x
        self.y = self.original_y

    def as_state(self) ->Tuple[int, int, int, int]:
        return(
            self.x,
            self.y,
            self.width,
            self.height
        )


# -*- coding: utf-8 -*-
"""
Created on Fri Jan 13 14:18:38 2017

@author: mguski
"""
# %%
import numpy as np
import matplotlib.pyplot as plt

nAntennas = 16
x_spacing = 15.4
freq = 8.1e6

nSteps = 360*5

c = 3e8
alphaVec = np.array(range(nSteps)) /nSteps*2* np.pi
# %%
import numpy as np

nBeams = 16
x_spacing = 15.4
beam_sep = 3.24
c = 3e8

center_beam = (nBeams -1)/2
delay_list = []

for iBeam in range(nBeams):
 azm = (iBeam - center_beam) * beam_sep /180 * np.pi
 delta_time = np.sin(azm) * x_spacing / c
 delay_list.append(delta_time*1e9)

print("Theory Values for:")
print("  {} beams, {} deg beam sep, {} m antenna spacing".format(nBeams, beam_sep,x_spacing))

for iBeam in range(nBeams):
    print("  beam {: >3d} / {} :   time diff {: >7.3f} ns".format(iBeam, nBeams, delay_list[iBeam]))

# %% how to map anteeas to PCs, if one PC fails

nAntennas = 16
x_spacing = 15.4
R = 1/nAntennas * np.sin(np.pi*x_spacing/c*freq*nAntennas * np.cos(alphaVec)) / np.sin(np.pi*x_spacing/c*freq*np.cos(alphaVec))
plt.polar(alphaVec, 20*np.log10(np.absolute(R))+100, label='normal')

nAntennas = 8
x_spacing = 15.4
R = 1/nAntennas * np.sin(np.pi*x_spacing/c*freq*nAntennas * np.cos(alphaVec)) / np.sin(np.pi*x_spacing/c*freq*np.cos(alphaVec))
plt.polar(alphaVec, 20*np.log10(np.absolute(R))+100, label='one side')

nAntennas = 8
x_spacing = 15.4*2
R = 1/nAntennas * np.sin(np.pi*x_spacing/c*freq*nAntennas * np.cos(alphaVec)) / np.sin(np.pi*x_spacing/c*freq*np.cos(alphaVec))
plt.polar(alphaVec, 20*np.log10(np.absolute(R))+100, label='every second')

plt.legend()


import bpy
from bpy.props import *
from ... events import executionCodeChanged
from ... sockets.info import toBaseDataType, isBase
from ... base_types import AnimationNode, ListTypeSelectorSocket

class SetListElementNode(AnimationNode, bpy.types.Node):
    bl_idname = "an_SetListElementNode"
    bl_label = "Set List Element"
    errorHandlingType = "MESSAGE"

    assignedType: ListTypeSelectorSocket.newProperty(default = "Float")

    clampIndex: BoolProperty(name = "Clamp Index", default = False,
        description = "Clamp the index between the lowest and highest possible index",
        update = executionCodeChanged)

    allowNegativeIndex: BoolProperty(name = "Allow Negative Index",
        description = "-2 means the second last list element",
        update = executionCodeChanged, default = False)

    def create(self):
        prop = ("assignedType", "BASE")
        self.newInput(ListTypeSelectorSocket(
            "List", "list", "LIST", prop, dataIsModified = True))
        self.newInput(ListTypeSelectorSocket(
            "Element", "element", "BASE", prop))
        self.newInput("Integer", "Index", "index")

        self.newOutput(ListTypeSelectorSocket(
            "List", "list", "LIST", prop))

    def drawAdvanced(self, layout):
        layout.prop(self, "clampIndex")
        layout.prop(self, "allowNegativeIndex")
        self.invokeSelector(layout, "DATA_TYPE", "assignListDataType",
            dataTypes = "LIST", text = "Change Type", icon = "TRIA_RIGHT")

    def getExecutionCode(self, required):
        if self.allowNegativeIndex:
            if self.clampIndex:
                yield "if len(list) != 0: list[min(max(index, -len(list)), len(list) - 1)] = element"
            else:
                yield "if -len(list) <= index < len(list): list[index] = element"
        else:
            if self.clampIndex:
                yield "if len(list) != 0: list[min(max(index, 0), len(list) - 1)] = element"
            else:
                yield "if 0 <= index < len(list): list[index] = element"
        yield "else: self.setErrorMessage('Index out of range')"

    def assignListDataType(self, listDataType):
        self.assignType(toBaseDataType(listDataType))

    def assignType(self, baseDataType):
        if not isBase(baseDataType): return
        if baseDataType == self.assignedType: return
        self.assignedType = baseDataType
        self.refresh()


'''
Created on 2012-7-2

@author: guoxc
'''
import json
import Tkinter
test_tree = json.load(open("root_frame.json"))
def tree_print(tr, intend=0):
    if not isinstance(tr, dict):
        raise Exception("tree not a dict")
    if "type" not in tr.keys() or "id" not in tr.keys():
        raise Exception("syntax Error")
    attr_list =  [(str(attr) +":"+str(tr[attr])) for attr in tr if attr not in ("menu","type","id","items")]
    print ("|   " * (intend - 1 if intend > 1 else 0) +
        "|___" * ( 1 if intend >= 1 else 0) + tr["type"] +
        ":" + tr["id"] + "  " +str(attr_list))
    if "menu" in tr.keys():
        tree_print(tr["menu"], intend +1)
    if "items" in tr.keys():
        for item in tr["items"]:
            tree_print(item, intend +1)

def get_menu(tr, parent):
    new_menu = Tkinter.Menu(parent, tearoff=0)
    if "items" in tr.keys():
        for item in tr["items"]:
            if item["type"] == "Menu":
                new_menu.add_cascade(menu=get_menu(item, new_menu),
                                     label=get_menu_text(item))
            elif item["type"] == "MenuItem":
                new_menu.add_command(label=item["text"], command=get_cmd_by_id(item["id"]))
    return new_menu

def get_button(tr, parent):
    opt = get_options(tr)
    new_button = Tkinter.Button(parent, **opt)
    return new_button

def get_entry(tr, parent):
    opt = get_options(tr)
    new_entry = Tkinter.Entry(parent, **opt)
    return new_entry

def get_label(tr, parent):
    opt = get_options(tr)
    new_label = Tkinter.Label(parent, **opt)
    return new_label

def get_options(opdict):
    opt = {}
    for key in opdict.keys():
        if key not in ["type", "id","position","items","class","sticky"]:
            opt[key] = opdict[key]
    return opt

def get_cmd_by_id(id_):
    if id_ +"_cmd" in globals():
        return globals()[id_ +"_cmd"]
    else:
        return None

def menu1_1_3_cmd():
    root.destroy()

def get_menu_text(tr):
    if not tr["text"]:
        return None
    return tr["text"]

def parse_position(text):
    rlist = text.split(" ")
    rlist = [int(item) for item in rlist]
    if len(rlist) == 2:
        return {"row": rlist[0],"column":rlist[1]}
    elif len(rlist) == 4:
        return {"row": rlist[0],"column":rlist[1],"rowspan":rlist[2],"columnspan":rlist[3]}

def tree_parse(tr):
    if tr["menu"]:
        root.config(menu=get_menu(tr["menu"], root))
    if tr["items"] and tr["layout"]:
        if tr["layout"] == "grid":
            for item in tr["items"]:
                position_dict = parse_position(item["position"])
                if "sticky" in item.keys():
                    position_dict["sticky"] = item["sticky"]
                if item["type"] == "Button":
                    get_button(item, root).grid(**position_dict)
                elif item["type"] == "Entry":
                    get_entry(item, root).grid(**position_dict)
                elif item["type"] == "Label":
                    get_label(item, root).grid(**position_dict)

root = Tkinter.Tk()
tree_print(test_tree)
tree_parse(test_tree)
root.mainloop()


from pathlib import Path
import os

def extract_video_id(url):
    url_split = url.split('?')
    if len(url_split) < 2:
        return null
    for param in url_split[1].split('&'):
        key, value = param.split("=")
        if key == "v":
            return value
    return null

def create_download_directory():
    """Tries to create download directory at startup."""

    path = Path(os.getcwd()).joinpath('download')
    try:
        Path(path).mkdir(parents=True, exist_ok=True)
    except (OSError,FileExistsError,FileNotFoundError) as error:
        print("Creation of downlaod directory failed")
        print(error)


"""The form for the settings."""

import curses
from general.functions import can_be_dir
from general.functions import acc_list_to_multiline
from general.functions import multiline_to_acc_list
from general.functions import move_list_entry
import npyscreen

class AfaListAction(npyscreen.MultiLineAction):
    """List holding the afa templates."""

    def __init__(self, *args, **kwargs):
        """Initialize the class."""
        super(AfaListAction, self).__init__(*args, **kwargs)

        # set up key shortcuts
        self.add_handlers({
            curses.KEY_IC: self.add_afa_entry,
            curses.KEY_DC: self.delete_afa_entry
        })

    def update_values(self):
        """Update values and refresh view."""
        self.values = self.parent.parentApp.S.get_afa_table()

        self.display()
        self.clear_filter()

    def add_afa_entry(self, keypress=None):
        """Add AfA entry."""
        # get the input from user
        name = npyscreen.notify_input(
            'Name of AfA entry:'
        )

        if not name:
            return

        account = npyscreen.notify_input(
            'Account string:',
            pre_text='afa'
        )

        if not account:
            return

        years = npyscreen.notify_input(
            'Years for AfA:'
        )

        if not years:
            return

        got_it = self.parent.parentApp.S.add_afa_table(
            name=name,
            account=account,
            years=years
        )

        if not got_it:
            npyscreen.notify_confirm(
                'Could not add the AfA table entry, sorry.',
                form_color='WARNING'
            )

        self.update_values()

    def delete_afa_entry(self, keypress=None):
        """Delete AfA entry."""
        # cancel if there are no values
        if len(self.values) < 1:
            return False

        selected = self.values[self.cursor_line]

        really = npyscreen.notify_yes_no(
            'Really delete template "{}"?'.format(selected['name'])
        )

        if really:
            deleted = self.parent.parentApp.S.del_afa_table(index=self.cursor_line)

            if not deleted:
                npyscreen.notify_confirm(
                    'Could not delete the entry, sorry!',
                    form_color='WARNING'
                )

            self.update_values()

    def display_value(self, vl):
        """Display the value."""
        return '[{}]'.format(vl['name'])

    def actionHighlighted(self, act_on_this, keypress):
        """Do something, because a key was pressed."""
        # cancel if there are no values
        if len(self.values) < 1:
            return False

        # get the input from user
        name = npyscreen.notify_input(
            'Name of AfA entry:',
            pre_text=act_on_this['name']
        )

        if not name:
            return

        account = npyscreen.notify_input(
            'Account string:',
            pre_text=act_on_this['account']
        )

        if not account:
            return

        years = npyscreen.notify_input(
            'Years for AfA:',
            pre_text=str(act_on_this['years'])
        )

        if not years:
            return

        got_it = self.parent.parentApp.S.add_afa_table(
            name=name,
            account=account,
            years=years
        )

        if not got_it:
            npyscreen.notify_confirm(
                'Could not modify the AfA table entry, sorry.',
                form_color='WARNING'
            )
        else:
            # delete the old entry
            deleted = self.parent.parentApp.S.del_afa_table(index=self.cursor_line)

            if not deleted:
                npyscreen.notify_confirm(
                    'Something went wrong, sorry!',
                    form_color='WARNING'
                )

        self.update_values()

class TitleAfaList(npyscreen.TitleMultiLine):
    """Inherit from TitleMultiLine, but get MultiLineAction."""

    _entry_type = AfaListAction

class PostingList(npyscreen.MultiLineEditable):
    """PostingList holding teh posts."""

    def __init__(self, *arg, **kwargs):
        """Initialize the class."""
        super(PostingList, self).__init__(*arg, **kwargs)

        # set up key shortcuts
        self.add_handlers({
            '+': self.move_up,
            '-': self.move_down
        })

        # set up additional multiline options
        self.scroll_exit = True

    def move_up(self, keypress=None):
        """Move selected offer up in the list."""
        lis = self.values

        # cancel if list is < 2
        if len(lis) < 2:
            return False

        # move selected item up
        new_index = move_list_entry(
            lis=lis,
            index=self.cursor_line,
            direction=1
        )

        self.cursor_line = new_index

    def move_down(self, keypress=None):
        """Move selected offer down in the list."""
        lis = self.values

        # cancel if list is < 2
        if len(lis) < 2:
            return False

        # move selected item up
        new_index = move_list_entry(
            lis=lis,
            index=self.cursor_line,
            direction=-1
        )

        self.cursor_line = new_index

class TitleMultiLineEdit(npyscreen.TitleText):
    """Titled MultiLineEdit"""

    _entry_type = PostingList

class SettingsForm(npyscreen.FormMultiPageActionWithMenus):
    """Form for the settings."""

    def __init__(self, *args, **kwargs):
        """Initialize the class."""
        super(SettingsForm, self).__init__(*args, **kwargs)

        # set up key shortcuts
        self.add_handlers({
            '^O': self.on_ok,
            '^Q': self.on_cancel
        })

    def add_afa(self):
        """Add afa table entry."""
        self.afa_table.entry_widget.add_afa_entry()

    def del_afa(self):
        """Delete afa table entry."""
        self.afa_table.entry_widget.delete_afa_entry()

    def exit(self):
        """Exit the programm."""
        self.parentApp.setNextForm(None)
        self.parentApp.switchFormNow()

    def create(self):
        """Create the form."""
        # create the menu
        self.m = self.new_menu(name='Menu')
        self.m.addItem(text='Add afa entry', onSelect=self.add_afa, shortcut='a')
        self.m.addItem(text='Del afa entry', onSelect=self.del_afa, shortcut='d')
        self.m.addItem(text='Exit', onSelect=self.exit, shortcut='e')

        # create input widgets
        self.data_path = self.add_widget_intelligent(
            npyscreen.TitleFilenameCombo,
            name='Data path:',
            begin_entry_at=26
        )
        self.ledger_path = self.add_widget_intelligent(
            npyscreen.TitleFilenameCombo,
            name='Ledger path:',
            begin_entry_at=26
        )
        self.def_state = self.add_widget_intelligent(
            npyscreen.TitleSelectOne,
            name='Default state:',
            begin_entry_at=26,
            values=['*', '!'],
            value=[0],
            slow_scroll=True,
            scroll_exit=True,
            max_height=2
        )
        self.def_code = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Default code:',
            begin_entry_at=26
        )
        self.def_payee = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Default payee:',
            begin_entry_at=26
        )
        self.def_commodity = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Default commodity:',
            begin_entry_at=26
        )
        self.def_postings = self.add_widget_intelligent(
            TitleMultiLineEdit,
            name='Default postings:',
            begin_entry_at=26,
            max_height=4
        )
        self.def_force_add = self.add_widget_intelligent(
            npyscreen.TitleMultiSelect,
            name='Force adding:',
            begin_entry_at=26,
            values=['enabled'],
            max_height=2,
            scroll_exit=True
        )
        self.dec_separator = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Decimal separator:',
            begin_entry_at=26
        )
        self.date_separator = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Date separator:',
            begin_entry_at=26
        )
        self.date_fmt = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Date format:',
            begin_entry_at=26
        )
        self.ledger_file = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='Ledger file:',
            begin_entry_at=26
        )
        self.split_years_to_files = self.add_widget_intelligent(
            npyscreen.TitleMultiSelect,
            name='Split years to files:',
            begin_entry_at=26,
            values=['enabled'],
            max_height=2,
            scroll_exit=True
        )
        self.afa_threshold_amount = self.add_widget_intelligent(
            npyscreen.TitleText,
            name='AfA threshold amount:',
            begin_entry_at=26
        )
        self.afa_table = self.add_widget_intelligent(
            TitleAfaList,
            name='AfA table:',
            begin_entry_at=26,
            scroll_exit=True
        )

    def beforeEditing(self):
        """Get values from settings object."""
        # change name and color according to the settings
        if self.parentApp.S._got_arguments:
            self.name = (
                'Ledgeradd > Settings (arguments altered settings - cannot save!'
            )
            self.color = 'DANGER'
        else:
            self.name = 'Ledgeradd > Settings'
            self.color = 'FORMDEFAULT'

        self.data_path.value = self.parentApp.S.data_path
        self.ledger_path.value = self.parentApp.S.ledger_path
        self.def_state.value = [self.def_state.values.index(self.parentApp.S.def_state)]
        self.def_code.value = self.parentApp.S.def_code
        self.def_payee.value = self.parentApp.S.def_payee
        self.def_commodity.value = self.parentApp.S.def_commodity

        self.def_postings.entry_widget.values = acc_list_to_multiline(
            self.parentApp.S.get_def_postings()
        )

        self.def_force_add.value = (
            [0] if self.parentApp.S.get_def_force_add() else []
        )

        self.dec_separator.value = self.parentApp.S.dec_separator
        self.date_separator.value = self.parentApp.S.date_separator
        self.date_fmt.value = self.parentApp.S.date_fmt
        self.ledger_file.value = self.parentApp.S.ledger_file
        self.split_years_to_files.value = (
            [0] if self.parentApp.S.get_split_years_to_files() else []
        )
        self.afa_threshold_amount.value = str(self.parentApp.S.get_afa_threshold_amount())
        self.afa_table.entry_widget.update_values()

    def on_ok(self, keypress=None):
        """Do something because user pressed ok."""
        # get values into temp variables
        data_path = self.data_path.value.rstrip('/')
        ledger_path = self.ledger_path.value.rstrip('/')

        # check correctness of values
        data_path_correct = can_be_dir(data_path)
        ledger_path_correct = can_be_dir(ledger_path)

        # dir is not creatable
        if not data_path_correct or not ledger_path_correct:
            message = npyscreen.notify_ok_cancel(
                'Data/ledger path is no valid folder name! Please change it!',
                title='Wrong folder names!',
                form_color='WARNING'
            )
            if message:
                # ok, I will edit them again
                self.editing = True
            else:
                # no, I will go back to main screen
                self.editing = False
                # swtich back
                self.parentApp.setNextForm('MAIN')
                self.parentApp.switchFormNow()

        # it is creatable, set new values to settings object
        else:
            # new values
            self.parentApp.S.data_path = data_path
            self.parentApp.S.ledger_path = ledger_path

            # defaults
            self.parentApp.S.def_state = self.def_state.values[self.def_state.value[0]]
            self.parentApp.S.def_code = self.def_code.value
            self.parentApp.S.def_payee = self.def_payee.value
            self.parentApp.S.def_commodity = self.def_commodity.value
            self.parentApp.S.set_def_postings(
                multiline_to_acc_list(self.def_postings.entry_widget.values)
            )
            self.parentApp.S.set_def_force_add(
                True if self.def_force_add.value == [0] else False
            )

            # formatting
            self.parentApp.S.dec_separator = self.dec_separator.value
            self.parentApp.S.date_separator = self.date_separator.value
            self.parentApp.S.date_fmt = self.date_fmt.value

            # file
            self.parentApp.S.ledger_file = self.ledger_file.value
            self.parentApp.S.set_split_years_to_files(
                True if self.split_years_to_files.value == [0] else False
            )

            # afa
            self.parentApp.S.set_afa_threshold_amount(self.afa_threshold_amount.value)
            self.parentApp.S.set_afa_table(self.afa_table.entry_widget.values)

            # store it
            self.parentApp.S.save_settings_to_file()

            # switch back
            self.parentApp.setNextForm('MAIN')
            self.parentApp.switchFormNow()

    def on_cancel(self, keypress=None):
        """Go back without changing a thing."""
        # swtich back
        self.parentApp.setNextForm('MAIN')
        self.parentApp.switchFormNow()


#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os, sys
from utils import measure

def execute(day):
  main = __import__(f'{day}.main').main.main
  os.chdir(day)
  measure(main)

def list_available_days():
  dirs = sorted([dir for dir in os.listdir(base_dir) if os.path.isdir(dir) and dir.isnumeric()])
  print('Available days:')
  i = 0
  while len(dirs) > i:
    print('\t'.join(dirs[i:i+5]))
    i += 5

if __name__ == '__main__':
  base_dir = os.path.realpath(os.path.dirname(__file__))
  os.chdir(base_dir)

  if len(sys.argv) >= 2 and sys.argv[-2] == '--day':
    execute(sys.argv[-1])
  else:
    list_available_days()
    day = input('> ')
    execute(day)


from datetime import datetime, timedelta

def check_time_restriction(start_time, end_time):
    st = start_time.split(':')
    et = end_time.split(':')

    message_time = datetime.now()
    start_time = message_time.replace(hour=int(st[0]), minute=int(st[1]), second=int(st[2]), microsecond=0)
    end_time = message_time.replace(hour=int(et[0]), minute=int(et[1]), second=int(et[2]), microsecond=0)

    # Tweak the date if the time range spans midnight
    if end_time.hour < start_time.hour:
        # If check time is after midnight, move start time back to the previous day
        if message_time.hour < end_time.hour:
            start_time -= timedelta(days=1)
        # .. and if before midnight, move end time to the next day.
        else:
            end_time += timedelta(days=1)

    if start_time <= message_time <= end_time:
        return True

    return False


from pathlib import Path
import pandas as pd

from prefect import flow,task
from prefect_gcp.cloud_storage import GcsBucket
from prefect_gcp import GcpCredentials

@task(log_prints=True,retries=2)
def extract_from_gcs(playlist_name:str) -> Path:
    """Download trip data from GCS"""

    gcs_path = f"data/DTC_{playlist_name}_playlist_snapshot.parquet"
    gcs_block = GcsBucket.load("zoom-gcs")
    gcs_block.get_directory(from_path=gcs_path, local_path=f"./data/")
    print(f"path: ./data/{gcs_path}")

    return Path(f"./data/{gcs_path}")

@task(log_prints=True)
def transform(path: Path) -> pd.DataFrame:
    """Data cleaning for bq"""
    print(path)
    df = pd.read_parquet(path)
    print(f"number of rows: {len(df)}")

    #df[''].fillna(df[''].mean(),inplace=True)

    return df

@task(log_prints=True,retries=2)
def write_bq(df: pd.DataFrame,playlist_name:str) -> None:
    """Write DataFrame to BigQuery"""
    gcp_credentials_block = GcpCredentials.load("gcp-credentials")
    df.to_gbq(destination_table=f"youtube_playlist.youtube_playlist_{playlist_name}_info",
        project_id="my-rides-antonis",
        credentials= gcp_credentials_block.get_credentials_from_service_account(),
        chunksize=500000,
        if_exists="replace"
    )
    return

@flow(name="ETL-TO-BIG-QUERY")
def etl_gcs_to_bq(playlist_name:str) -> None:
    """Main ETL flow to load data into Big Query"""

    path = extract_from_gcs(playlist_name)
    df = transform(path)
    write_bq(df,playlist_name)

@flow(name="Parent-ETL-TO-BIG-QUERY")
def etl_parent_gcs_to_bq(playlist_name_list:list[str]) -> None:
    """Parent of ETL flow to load data into Big Query"""

    #for every playlist create a different table
    for playlist_name in playlist_name_list:
        etl_gcs_to_bq(playlist_name)

if __name__ == '__main__':
    playlist_name_list = ['Data Engineering Zoomcamp 2023','Machine Learning Zoomcamp 2022','MLOps Zoomcamp 2022']
    etl_parent_gcs_to_bq(playlist_name_list)


import numpy as np
from sklearn import neighbors, datasets, tree, linear_model
from sklearn.model_selection import cross_val_score
from sklearn.externals import joblib

import classsol

#load input data
words = []
with open("words.txt", encoding="latin-1") as file:
    for line in file:
        line = line.split(' ') #or some other preprocessing
        words.append(line) #storing everything in memory!

X = words[0]

for test in ["wordsclass.npy", "wordsclass2.npy"]:
    print("Testing " + test)
    #load output data
    Y=np.load(test)

    f = classsol.features(X)

    # classsol.test_params(f, Y)

    clf = classsol.mytraining(f,Y)

    Ypred = classsol.myprediction(f, clf)

    if (np.sum(Y^Ypred)/len(X))<.05:
        print("Erro bastante baixo. PERFECT!\n")
    elif (np.sum(Y^Ypred)/len(X))<.3:
        print("Erro nos Q dentro dos limites de tolerância. OK\n")
    else:
        print("Erro nos Q acima dos limites de tolerância. FAILED\n")


import requests
from decouple import config
import os
#Getting configs
url = config('url')
min_value = int(config('min_value'))
max_value = int(config('max_value'))
file_format = config('file_format')
num_pad = int(config('num_pad'))

#create downloads folder
if not os.path.isdir('downloads'):
    os.mkdir('downloads')

for current_value in range(min_value, max_value+1):

    if(num_pad):
        current_value = str(current_value).zfill(num_pad)
    r = requests.get(url.format(current_value))
    open('downloads/{}.{}'.format(current_value, file_format), 'wb').write(r.content)
    print("downloaded {}.{}".format(current_value, file_format))


# print the largest palindrome possible from two numbers with 3 digits

result=[]
for i in range(100,1000):
    for j in range(100,1000):
        if (i*j)==(j*i):
            result.append(i*j)
print(max(result))

# printing all 3 digit palindrome in python
result=[]
for i in range(100,1000):
    temp=i
    rev=0
    while(i>0):
        dig=i%10
        rev=rev*10+dig
        i=i//10
    if(temp==rev):
        result.append(temp)
print(result)


import math

class Accion:
    DERECHA = 1
    IZQUIERDA = 0
    #NADA = 2
    maxValor = 1

class Posicion:
    BORDE_IZQUIERDO = 1
    BORDE_DERECHO = 2
    CENTRO = 3

class Estado:
    angulo = 0
    velocidadAngular = 0
    posicion = Posicion.CENTRO

    def __init__(self):
        self.angulo = math.pi
        self.velocidadAngular = 0
        posicion = 0

    def __repr__(self):
        return "Angulo: {}, Velocidad Angular: {}, Posicion: {})".format(
            self.angulo,
            self.velocidadAngular,
            self.posicion)


import torch

import warnings
from scipy.spatial.distance import cdist

# supress UserWarning
warnings.filterwarnings("ignore", category=UserWarning)
from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
model = PretrainedSpeakerEmbedding(
    "speechbrain/spkrec-ecapa-voxceleb",
    device="cuda" if torch.cuda.is_available() else "cpu",)

from pyannote.audio import Audio

audio = Audio(sample_rate=16000, mono="downmix")

waveform1, sample_rate = audio({"audio": "../taylor-swift/data/vocals30.wav"})
embedding1 = model(waveform1[None])

import os

all_sims = []

for file in os.listdir("../taylor-swift/data/"):
    if file.endswith(".wav") and "_m" not in file:
        waveform2, sample_rate = audio({"audio": "../taylor-swift/data/" + file})
        embedding2 = model(waveform2[None])

        # compare embeddings using "cosine" distance
        distance = cdist(embedding1, embedding2, metric="cosine")

        all_sims.append(distance[0][0])

avg = sum(all_sims) / len(all_sims)

print(avg)


import csv

import sys

import requests

from bs4 import BeautifulSoup as bs

import pandas as pd

filename = "gamelog.csv"

TAGS = {"PTS":26, "BLK": 23, "STL": 22, "AST": 21, "TRB": 20}

def getInfo(games, stat):
    global TAGS

    data = []

    stat = TAGS[stat]

    total_number = len(games)

    for i in games:

        data.append(i)
    sum = 0

    for i in range(len(data)):

        sum += float(data[i][stat])

        data[i] = float(data[i][stat])

    data = sorted(data)

    mean = sum/total_number

    if total_number % 2 == 1: #odd

        median = data[int(total_number/2)]

    else:
        median = (data[int(total_number/2)]+data[int(1+total_number/2)])/2
    #next we want to get the varience and the numbers above certain thresholds
    above = [0, 0, 0, 0, 0, 0, 0] #5-35
    var = 0
    max = data[0]
    min = data[0]
    for i in data:
        var += (i-mean)**2
        if i > max:
            max = i
        if i < min:
            min = i
        if i/5 >= 1:
            if int(i/5) > 7:
                goto = 7
            else:
                goto = int(i/5)
            for j in range(goto):
                above[j] += 1
    var *= (1/(total_number-1))

    print("Mean: ", mean, "Median: ", median, "Varience: ", var, "\n")
    print("Min: ", min, "Max: ", max)
    for i in range(len(above)):
        try:
            odds = "%0.3f" % float(total_number/above[i])
        except Exception as e:

            odds = "0"

        print(str(5*(i+1)) + "+: " + str(above[i]) + "\t" + str(int(100*above[i]/total_number)) + "%\tMin Odds: " +odds )

def openfile(filename):

    fields = []
    rows = []
    with open(filename, 'r') as csvfile:
        # creating a csv reader object
        csvreader = csv.reader(csvfile)

        # extracting field names through first row
        fields = next(csvreader)

        # extracting each data row one by one
        for row in csvreader:
            rows.append(row)

        # get total number of rows

    #get rid of the blank rows

    newrows = []
    for i in range(len(rows)):

        if rows[i][1] != '':
            newrows.append(rows[i])

    return fields, newrows

def getFile(link, year):
    print("Getting Data From ", year)
    link = "https://www.basketball-reference.com/players/" + link[0] +"/" +link + "/gamelog/" + year

    r = requests.get(link)

    soup = bs(r.content, 'html.parser')

    table = soup.find('table', id='pgl_basic')
    def rowgetDataText(tr, coltag='td'): # td (data) or th (header)
        return [td.get_text(strip=True) for td in tr.find_all(coltag)]
    rows = []

    trs = table.find_all('tr')
    headerow = rowgetDataText(trs[0], 'th')
    if headerow: # if there is a header row include first
        rows.append(headerow)
        trs = trs[1:]
    for tr in trs: # for every table row
        rows.append(rowgetDataText(tr, 'td') ) # data row
    new_table = []

    removeLines = ['Did Not Dress', 'Inactive', 'Did Not Play', 'Not With Team', 'Player Suspended']

    for i in rows:

        if i != [] and i[7] not in removeLines and i[0] != 'Rk':
            new_table.append(i)

    #from here we have a table which just needs to be converted to csvfile

    csvfile = ""

    for i in new_table:

        for j in i:

            csvfile += j
            csvfile +=','

        csvfile+='\n'
    return csvfile

def writeCSV (csvText, mode):

    if mode == "w":

        with open('gamelog.csv', 'w') as out:
            out.write(csvText)

    elif mode == "a":

        with open('gamelog.csv', 'a') as out:
            out.write(csvText)

def lookup(name):

    link = ""

    r = requests.get("https://www.basketball-reference.com/players/" + name.split()[1][0].lower() + "/")

    soup = bs(r.content, 'html.parser')

    link = "https://www.basketball-reference.com" + soup.find("a", string=name).get("href").replace(".html", "/gamelog")

    return soup.find("a", string=name).get("href").split("/")[3].split(".")[0]

def getSeasons(name):

    seasons = []

    r = requests.get("https://www.basketball-reference.com/players/" + name[0] + "/"+ name + ".html")

    soup = bs(r.content, 'html.parser')

    table = soup.find("table", id="per_game")

    tbody = table.find('tbody')
    th_head = tbody.find_all('th')

    for i in th_head:

        seasons.append(i.text)

    #this gets rid of duplicates
    seasons = list(dict.fromkeys(seasons))

    #this cuts it down to only the first yea rin the season ex 2015-16 -> 2015

    for i in range(len(seasons)):

        seasons[i] = str(seasons[i][:2])+str(seasons[i][5:])

    return seasons

playerSeasons = getSeasons(lookup(sys.argv[1]))

print("Saving Gamelogs")

writeCSV(getFile(lookup(sys.argv[1]), playerSeasons[0]), "w")

for i in range(1, len(playerSeasons)):

    writeCSV(getFile(lookup(sys.argv[1]), playerSeasons[i]), "a")

ntn, smt = openfile(filename)

getInfo(smt, sys.argv[2])


import csv
import math
import tensorflow as tf
import random

class CongnitiveNavigationDatasnap:
  def __init__(obj):
    obj.Row = []
    obj.Input = []
    obj.Output = []

def mount_training_batch(dataset,batch_size):
  selected_numbers = []
  new_dataset = CongnitiveNavigationDatasnap()

  while len(selected_numbers) < batch_size:
    index = random.randint(0,(len(dataset.Input)-1))
    if(not(index in selected_numbers)):
      new_dataset.Input.append(dataset.Input[index])
      new_dataset.Output.append(dataset.Output[index])
      selected_numbers.append(index)

  return new_dataset

def write_csv_row(csv_row,csv_name):
  with open((csv_name+'.csv'),'a') as fd:
    writer = csv.writer(fd)
    writer.writerow(csv_row)
    fd.close()

def generate_X_tensor(ESE45, ESE30, ESF, ESD30, ESD45, EAR):
  ESE45 = float(ESE45)
  ESE30 = float(ESE30)
  ESF = float(ESF)
  ESD30 = float(ESD30)
  ESD45 = float(ESD45)
  EAR = round(float((EAR)/(3.1415)), 2)

  X =[ [ESD30,ESD45,ESF, EAR],
      [ESD45,ESF,ESE45, EAR],
      [ESE30,ESE45, ESF, EAR]
    ]

  return X

def generate_Y_tensor(SVL, SVA):
  # SVL = round(((float(SVL)+1)/2),2)
  # SVA = round(((float(SVA)+1)/2),2)

  Y = [float(SVL),float(SVA)]

  return Y

def load_dataset(file):

    dataset = CongnitiveNavigationDatasnap()

    with open(file) as csvfile:
        spamreader = csv.reader(csvfile, delimiter=",")
        for row in spamreader:
            new_input = generate_X_tensor(float(row[0]),float(row[1]),float(row[2]),float(row[3]),float(row[4]), float(row[5]))
            new_output = generate_Y_tensor(row[6], row[7])
            dataset.Input.append(new_input)
            dataset.Output.append(new_output)

        return dataset

def get_network_model():
  model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(3, 4)),
    tf.keras.layers.Dense(24, activation='tanh'),
    tf.keras.layers.Dense(18, activation='tanh'),
    tf.keras.layers.Dense(2, activation='tanh')
    ])
  # model = tf.keras.models.Sequential()
  # model.add(tf.keras.layers.InputLayer(input_shape =(3,4)))
  # model.add(tf.keras.layers.Dense(6,input_shape =(3,4),activation="tanh"))
  # model.add(tf.keras.layers.Dense(12,input_shape =(3,6),activation="tanh"))
  # model.add(tf.keras.layers.Dense(24,input_shape =(3,12), activation="tanh"))
  # model.add(tf.keras.layers.Dense(12,input_shape =(3,24), activation="tanh"))
  # model.add(tf.keras.layers.Dense(9,input_shape =(3,12), activation="tanh"))
  # model.add(tf.keras.layers.Flatten(input_shape=(3, 9)))
  # model.add(tf.keras.layers.Dense(27, activation="tanh"))
  # model.add(tf.keras.layers.Dense(12, activation="tanh"))
  # model.add(tf.keras.layers.Dense(9, activation="tanh"))
  # model.add(tf.keras.layers.Dense(2, activation="tanh"))

  model.compile(optimizer='adam',
              loss="mean_squared_error",
              metrics=['accuracy'])
  return model


import sys
import traceback

import network
import utils

IS_MISSING = -1

def approxEqual(value, target, tolerance):
   if (abs(target) <= tolerance ): return abs(value) <= tolerance
   return abs(float(value) / target - 1) <= tolerance

def check(name, value, target, tolerance):
   if approxEqual(value, target, tolerance):
      return True
   else:
      print("\nWrong %s: your value %f, correct value %f"
               % (name, value, target))
      return False

def checkExact(name, value, target):
   if value == target:
      return True
   else:
      print("\nWrong %s: your value %s, correct value %s"
               % (name, value, target))
      return False

def TDSP(testFileName):

   print("Running TDSP test: " + str(testFileName) + "...", end='')

   try:
      with open(testFileName, "r") as testFile:
         # Read test information
         try:
            fileLines = testFile.read().splitlines()
            pointsPossible = IS_MISSING
            networkFile = IS_MISSING
            origin = IS_MISSING
            departureTime = IS_MISSING
            correctCost = IS_MISSING
            travelTimes = dict()
            correctBacklink = IS_MISSING
            for line in fileLines:
               # Ignore comments and blank lines
               if len(line.strip()) == 0 or line[0] == '#':
                  continue

               # Set points possible
               if pointsPossible == IS_MISSING:
                  pointsPossible = int(line)
                  continue

               # Set network file
               if networkFile == IS_MISSING:
                  networkFile = line
                  testNetwork = network.Network(networkFile)
                  continue

               # Set origin/departure time
               if origin == IS_MISSING:
                  inputs = line.split(",")
                  origin = int(inputs[0])
                  departureTime = int(inputs[1])
                  continue

               # Set time-dependent travel times
               if len(travelTimes) < testNetwork.numLinks:
                  inputs = line.split(",")
                  testNetwork.links[inputs[0]].travelTime = [int(x) for x in inputs[1:]]
                  travelTimes[inputs[0]]  = True
                  continue

               # Set correct cost labels
               if correctCost == IS_MISSING:
                  inputs = line.split(",")
                  correctCost = [int(x) for x in inputs]
                  continue

               # Set correct backlink labels
               if correctBacklink == IS_MISSING:
                  inputs = line.split(",")
                  correctBacklink = inputs
                  continue

         except:
            print("\nError running test %s, attempting to continue with remaining tests.  Exception details: " % testFileName)
            traceback.print_exc(file=sys.stdout)
            return 0, 0

         # Now run the actual test
         try:
            (studentCost, studentBacklink) = testNetwork.TDSP(origin - 1, departureTime)
            numCorrect = 0
            for i in range(testNetwork.numNodes):
               numCorrect += 1 if check("Node %d cost" % (i+1), studentCost[i], correctCost[i], 0.01)  else 0
               numCorrect += 1 if checkExact("Node %d backlink" % (i+1), studentBacklink[i], correctBacklink[i]) else 0
         except utils.NotYetAttemptedException:
            print("...not yet attempted")
            return 0, pointsPossible
         except:
            print("\nException raised, attempting to continue:")
            traceback.print_exc(file=sys.stdout)
            print("\n...fail")
            return 0, pointsPossible

         if numCorrect < 2 * testNetwork.numNodes:
            print("...fail")
         else:
            print("...pass")
         return pointsPossible * (numCorrect / float(2 * testNetwork.numNodes) ), pointsPossible

   except IOError:
      print("\nError running test %s, attempting to continue with remaining tests.  Exception details: " % testFileName)
      traceback.print_exc(file=sys.stdout)
      return 0, 0

def convexCombo(testFileName):

   print("Running convex combinations test: " + str(testFileName) + "...", end='')

   try:
      with open(testFileName, "r") as testFile:
         # Read test information
         try:
            fileLines = testFile.read().splitlines()
            pointsPossible = IS_MISSING
            networkFile = IS_MISSING
            stepSize = IS_MISSING
            originalPathFlows = dict()
            targetPathFlows = dict()
            correctPathFlows = dict()
            for line in fileLines:
               # Ignore comments and blank lines
               if len(line.strip()) == 0 or line[0] == '#':
                  continue

               # Set points possible
               if pointsPossible == IS_MISSING:
                  pointsPossible = int(line)
                  continue

               # Set network file
               if networkFile == IS_MISSING:
                  networkFile = line
                  testNetwork = network.Network(networkFile)
                  continue

               # Set step size
               if stepSize == IS_MISSING:
                  stepSize = float(line)
                  continue

               # Set original path flows
               if len(originalPathFlows) < sum(len(OD.paths) for OD in testNetwork.ODs):
                  inputs = line.split(",")
                  path = tuple(inputs[:-testNetwork.timeHorizon])
                  flows = inputs[-testNetwork.timeHorizon:]
                  testNetwork.pathFlows[path] = [float(x) for x in flows]
                  originalPathFlows[path]  = True
                  continue

               # Set target path flows
               if len(targetPathFlows) < sum(len(OD.paths) for OD in testNetwork.ODs):
                  inputs = line.split(",")
                  path = tuple(inputs[:-testNetwork.timeHorizon])
                  flows = inputs[-testNetwork.timeHorizon:]
                  targetPathFlows[path] = [float(x) for x in flows]
                  continue

               # Set correct path flows
               if len(correctPathFlows) < sum(len(OD.paths) for OD in testNetwork.ODs):
                  inputs = line.split(",")
                  path = tuple(inputs[:-testNetwork.timeHorizon])
                  flows = inputs[-testNetwork.timeHorizon:]
                  correctPathFlows[path] = [float(x) for x in flows]
                  continue

         except:
            print("\nError running test %s, attempting to continue with remaining tests.  Exception details: " % testFileName)
            traceback.print_exc(file=sys.stdout)
            return 0, 0

         # Now run the actual test
         try:
            testNetwork.updatePathFlows(targetPathFlows, stepSize)
            numCorrect = 0
            for path in targetPathFlows:
               for t in range(testNetwork.timeHorizon):
                  numCorrect += 1 if check("Path %s flow at time %d" % (path, t), testNetwork.pathFlows[path][t], correctPathFlows[path][t], 0.01)  else 0
         except utils.NotYetAttemptedException:
            print("...not yet attempted")
            return 0, pointsPossible
         except:
            print("\nException raised, attempting to continue:")
            traceback.print_exc(file=sys.stdout)
            print("\n...fail")
            return 0, pointsPossible

         if numCorrect < len(targetPathFlows) * testNetwork.timeHorizon:
            print("...fail")
         else:
            print("...pass")
         return pointsPossible * (numCorrect / float(len(targetPathFlows) * testNetwork.timeHorizon) ), pointsPossible

   except IOError:
      print("\nError running test %s, attempting to continue with remaining tests.  Exception details: " % testFileName)
      traceback.print_exc(file=sys.stdout)
      return 0, 0


import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from batterycontroller import BatteryController
from central_planner import solve_bat

def get_cost(load, price_buy, price_sell):
    cost = 0
    for l, pb, ps in zip(load, price_buy, price_sell):
        cost += l * pb if l > 0 else l * ps
    return cost

def plot_results(df, day):

    df_ = df[df.date == day].copy()

    net_load = df_[df.type == "load"].iloc[:, 4:].values.sum(axis=0)
    net_batt = df_[df.type == "batt"].iloc[:, 4:].values.sum(axis=0)
    net_cng = df_[df.type == "central_gn"].iloc[:, 4:].values.sum(axis=0)
    net_cgl = df_[df.type == "central_gl"].iloc[:, 4:].values.sum(axis=0)
    net_cgb = df_[df.type == "central_gb"].iloc[:, 4:].values.sum(axis=0)
    fig, ax = plt.subplots()
    ax.plot(net_load, label="Net load {0}".format(np.abs(net_load).sum()))
    ax.plot(net_batt, label="Net batt {0}".format(np.abs(net_batt).sum()))
    ax.plot(net_cng, label="Net cng {0}".format(np.abs(net_cng).sum()))
    ax.plot(net_cgl, label="Net cgl {0}".format(np.abs(net_cgl).sum()))
    ax.plot(net_cgb, label="Net cgb {0}".format(np.abs(net_cgb).sum()))

    ax.legend()

def run_exp():

    ## Config

    SUN_TIMES = (15, 30)
    randomstate = np.random.RandomState(1234)
    T = 48

    price_1 = np.ones(T) * 12.0
    price_1[T // 2 :] = 16.0
    price_2 = np.ones(T) * 14.0
    price_s = np.ones(T) * 10.0

    b_max = 13
    b_min = 0
    e_c = 0.95
    e_d = 0.95
    d_max = 3
    d_min = -3
    battery = BatteryController(0, b_max, b_min, e_c, e_d, d_max, d_min)

    ## Reading data

    df = pd.read_csv("customers_data.csv")
    df["day"] = df.date.map(lambda x: x[:10])

    days_list = df.day.unique()
    users_list = df.customer.unique()
    N = len(users_list)

    ## Processing

    dataset = []

    start = time.time()

    for day in days_list[:5]:

        loads_ = []
        pss_ = []
        pbs_ = []
        costs_load_ = []
        costs_batt_ = []
        for i, user in enumerate(users_list[:]):

            load = df[(df.customer == user) & (df.day == day)].power.values.copy()
            solar = np.zeros_like(load)
            if i < N // 2:
                solar[SUN_TIMES[0] : SUN_TIMES[1]] = randomstate.uniform(
                    -1, 0, SUN_TIMES[1] - SUN_TIMES[0]
                )

            load += solar

            if (i > N // 4) and (i < N // 2):
                pb = price_2
            else:
                pb = price_1
            ps = price_s

            status, obj_bat, sol_bat = battery.find_optimal_step(load, pb, ps)
            if status != "optimal":
                raise ValueError("Optimization problem not optimal")

            cost_load = get_cost(load, pb, price_s)
            dataset.append((day, user, "load", cost_load) + tuple(load))
            dataset.append((day, user, "batt", obj_bat) + tuple(sol_bat + load))

            # Data for central planner
            loads_.append(load.copy())
            pbs_.append(pb.copy())
            pss_.append(ps.copy())
            costs_load_.append(cost_load)
            costs_batt_.append(obj_bat)

        status, obj, sol, om = solve_bat(
            loads_, pbs_, pss_, costs_load_, 0, 13, 0, 0.95, 0.95, 3, -3
        )
        if status != "optimal":
            raise ValueError("The big problem in day {0} broke".format(day))

        for i, s in enumerate(sol):
            net = (s + loads_[i]).copy()
            cost_ = get_cost(net, pbs_[i], pss_[i])
            dataset.append((day, users_list[i], "central_gl", cost_) + tuple(net))

        status, obj, sol, om = solve_bat(
            loads_, pbs_, pss_, costs_batt_, 0, 13, 0, 0.95, 0.95, 3, -3
        )
        if status != "optimal":
            raise ValueError("The big problem in day {0} broke".format(day))

        for i, s in enumerate(sol):
            net = (s + loads_[i]).copy()
            cost_ = get_cost(net, pbs_[i], pss_[i])
            dataset.append((day, users_list[i], "central_gb", cost_) + tuple(net))

        status, obj, sol, om = solve_bat(
            loads_, pbs_, pss_, [None] * len(loads_), 0, 13, 0, 0.95, 0.95, 3, -3
        )
        if status != "optimal":
            raise ValueError("The big problem in day {0} broke".format(day))

        for i, s in enumerate(sol):
            net = (s + loads_[i]).copy()
            cost_ = get_cost(net, pbs_[i], pss_[i])
            dataset.append((day, users_list[i], "central_gn", cost_) + tuple(net))

    columns = ["date", "user", "type", "cost"] + ["var{0}".format(x) for x in range(48)]
    final_dataset = pd.DataFrame(dataset)
    final_dataset.columns = columns

    elapsed = time.time() - start
    print("Elapsed time: {0}".format(elapsed))

    final_dataset.to_csv("final_dataset.csv", index=False)

if __name__ == "__main__":
    run_exp()


# https://leetcode.com/problems/middle-of-the-linked-list/

print('__name__:', __name__)

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

class ListOperation:
    def middle_node(self, head):
        """Finds out one middle node

            Retrieves the middle node in the linked list, if there are 2 middle
            nodes, return the second one.

            Args:
                head: The node points to the linked list.
                slow: The pointer traverse the linked list.
                fast: The pointer traverse as fast as twice as the 'slow' pointer.

            Returns:
                slow: 'slow' points to the middle node in the linked list finally.

            Analysis:
                Focus on better time complexity, we think of O(1), O(logN) and O(N).
                Since it asks us to find out the middle node in a linked
                list, we consider O(N).
                Use two pointers 'fast' and 'slow' both point to head initially.
                The 'fast' traverse as fast as twice than 'slow'.
                        When 'fast' reaches the end of list, 'slow' is in the middle.
                E.g. 1-2-3-4-5 returns 3,  1-2-3-4-5-6 returns 4

            Pseudocode:
                fast, slow = head, head
                if not head: return None
                while fast.next
                    if not fast.next.next
                        return slow.next
                    fast = fast.next.next
                    slow = slow.next
                return slow
        """
        # Code
        # Corner case though its a non-empty list
        if not head:
            return None

        fast = slow = head  # The 'fast' and 'slow' pointers

        while fast and fast.next:
            fast = fast.next.next
            slow = slow.next

        return slow

        # Follow up:
        # Input: [1, 2, 3, 4, 5, 6]
        # Output: (Serialization: [4, 5, 6])
        # Follow up Aanswer:
        # while slow:
        #     print(slow.val, end=',')
        #     slow = slow.next

list_operation = ListOperation()

def test_middle_node(lists):

    for i, llist in enumerate(lists):
        curr = head = None
        for j in llist:
            if head is not None:
                curr.next = ListNode(j)
                curr = curr.next
            else:
                head = ListNode(j)
                curr = head
        print('TC{}:\t'.format(i+1), list_operation.middle_node(head))

if __name__ == "__main__":
    print('__name__:', "執行此檔案才會顯示")

    lists = [
        [1, 2, 3, 4, 5],
        [1, 2, 3, 4, 5, 6],
        [],
        [1],
        [1, 2],
        [1, 2, 3],
        [1, 2, 3, 4],
        [1, 2, 3, 4, 5],
        [6, 5, 4, 3, 2, 1],
        [1, 0, 5, 7, 2, 4, 3, 15, 8, 12],
    ]
    test_middle_node(lists)

1-Easy/876. Middle of the Linked List/aaa.py


"""
Code for populating table 3

Run this in a directory that contains the relevant spectrum text files
"""
import os
import numpy as np
import pyspeckit
import goddi_nh3_fits
from hf_only_model import hfonly_fitter, sixsix_movinghf_fitter
from astropy import units as u
from astropy import constants
from astroquery.splatalogue import Splatalogue
import pylab as pl

# we're going to make a lot of figures.  Close some.
for ii in pl.get_fignums():
    pl.close(ii)

savepath = '/Users/adam/Dropbox/w51_goddi_share/spec-casa/figures/'
resultpath = '/Users/adam/Dropbox/w51_goddi_share/fit_results.txt'

outf = open(resultpath, 'w')

distance = 5.41*u.kpc
XNH3 = 1e-7

areas = {'e2e': np.pi*0.45*u.arcsec*0.4459*u.arcsec,
         'e2abs': np.pi*0.45*u.arcsec*0.40*u.arcsec,
         'e8': np.pi*2.1707*u.arcsec*1.6168*u.arcsec,
         'e2nw': np.pi*0.450*u.arcsec*0.400*u.arcsec,
         'e2clump': (10.87-0.58)*u.arcsec**2,
        }

freqs = {'66': 25.05596*u.GHz,
         '77': 25.71514*u.GHz,
         '99': 27.47794*u.GHz,
         '1010': 28.60474*u.GHz,
         '1313': 33.15685*u.GHz,
        }
eupper = {'66': 409,
          '77': 539,
          '99': 853,
          '1010': 1037,
          '1313': 1693,
         }
degeneracy = {'66': 2*(2*6+1),
              '77': 1*(2*7+1),
              '99': 2*(2*9+1),
              '1010': 1*(2*10+1),
              '1313': 1*(2*13+1),
             }

nh3tables = {line:Splatalogue.query_lines(freqs[line]*0.99999,
                                          freqs[line]*1.00001,
                                          line_strengths=['ls1','ls2','ls3','ls4','ls5'],
                                          chemical_name=' NH3 ')
             for line in freqs}
linestrengths = {'66':float(nh3tables['66'][nh3tables['66']['Linelist']=='JPL']['Sijμ2 (D2)'])*1e-18*u.esu*u.cm,
                 '77':float(nh3tables['77'][nh3tables['77']['Linelist']=='JPL']['Sijμ2 (D2)'])*1e-18*u.esu*u.cm,
                 '99':float(nh3tables['99'][nh3tables['99']['Linelist']=='JPL']['Sijμ2 (D2)'])*1e-18*u.esu*u.cm,
                 '1010':float(nh3tables['1010'][nh3tables['1010']['Linelist']=='JPL']['Sijμ2 (D2)'])*1e-18*u.esu*u.cm,
                 '1313':float(nh3tables['1313'][nh3tables['1313']['Linelist']=='JPL']['Sijμ2 (D2)'])*1e-18*u.esu*u.cm,
                }

FWHM = np.sqrt(8*np.log(2))

def tbl_vals(pi):
    """
    Compute the average amplitude, average width, and respective errors from the hyperfine fits
    """
    amp = (pi[2] + pi[5])/2.
    amperr = (pi[2].error + pi[5].error)/2.
    sigma = (pi[3] + pi[6])/2.
    sigmaerr = (pi[3].error + pi[6].error)/2.
    integral = amp * (2*np.pi)**0.5 *(sigma)
    int_error = (((amperr/amp)**2 + (sigmaerr/sigma)**2)*integral**2)**0.5

    return amp, amperr, sigma*FWHM, sigmaerr*FWHM, integral, int_error, pi[0].value, pi[0].error

def integral_hf(pi):
    return tbl_vals(pi)[4:6]

def tbl_vals_gaussian(pi):
    """
    Compute the average amplitude, average width, and respective errors from the hyperfine fits
    """
    amp = pi[0].value
    amperr = pi[0].error
    sigma = pi[2].value
    sigmaerr = pi[2].error
    integral = amp * (2*np.pi)**0.5 *(sigma)
    int_error = (((amperr/amp)**2 + (sigmaerr/sigma)**2)*integral**2)**0.5

    return amp, amperr, sigma*FWHM, sigmaerr*FWHM, integral, int_error, pi[1].value, pi[1].error

def integral_ml(pi):
    return tbl_vals_gaussian(pi)[4:6]

F = False
T = True

sp6e8,spK6e8,_,_ = goddi_nh3_fits.load_spectrum(6, object='w51e8', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp6e8.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp7e8,spK7e8,_,_ = goddi_nh3_fits.load_spectrum(7, object='w51e8', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp7e8.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp9e8,spK9e8,_,_ = goddi_nh3_fits.load_spectrum(9, object='w51e8', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp9e8.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp10e8,spK10e8,_,_ = goddi_nh3_fits.load_spectrum(10, object='w51e8', headerfile='/Users/adam/work/w51/goddi/W51-29GHzcont.map.image.fits')
sp10e8.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp13e8,spK13e8,_,_ = goddi_nh3_fits.load_spectrum(13, object='w51e8', headerfile='/Users/adam/work/w51/goddi/W51-33GHzcont.map.image.fits')
sp13e8.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)

sp6e8.plotter()
sp7e8.plotter()
sp9e8.plotter()
sp10e8.plotter()
sp13e8.plotter()
limits = [(50, 65), (25,30), (0, 1), (0.1, 6), (30, 35), (0, 1), (0.1, 6)]
limited = [(True,True)]*7
tied = ['']*7
tied[6] = 'p[3]'
tied[5] = 'p[2]'
sp6e8.specfit(fittype='hfonly', guesses=[58, 26.9, 0.02, 2.0, 31.4, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp7e8.specfit(fittype='hfonly', guesses=[58, 27.3, 0.02, 2.0, 31.2, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp9e8.specfit(fittype='hfonly', guesses=[58, 27.0, 0.02, 2.0, 30.1, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp6e8.specfit.plot_fit(show_hyperfine_components=True)
sp7e8.specfit.plot_fit(show_hyperfine_components=True)
sp9e8.specfit.plot_fit(show_hyperfine_components=True)
hfpi6e8 = sp6e8.specfit.parinfo
hfpi7e8 = sp7e8.specfit.parinfo
hfpi9e8 = sp9e8.specfit.parinfo
print("e8", file=outf)
print("hf 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp6e8.specfit.parinfo))), file=outf)
print("hf 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp7e8.specfit.parinfo))), file=outf)
"""
6-6 peak=0.0552312979281 +/- 0.00339213398418, fwhm=12.0092504285 +/- 0.346442710066, integral=3.60073813128 +/- 0.265490554173, center=59.8513485763 +/- 0.165718109489
7-7 peak=0.0343580596483 +/- 0.00211133669018, fwhm=7.08292886544 +/- 0.477361456933, integral=0.779163995174 +/- 0.0883610449159, center=59.9800528242 +/- 0.29151363486
"""
sp6e8.plotter.savefig(os.path.join(savepath, 'sp6e8_hf.png'))
sp7e8.plotter.savefig(os.path.join(savepath, 'sp7e8_hf.png'))
sp9e8.plotter.savefig(os.path.join(savepath, 'sp9e8_hf.png'))
sp10e8.plotter.savefig(os.path.join(savepath, 'sp10e8_hf.png'))
sp13e8.plotter.savefig(os.path.join(savepath, 'sp13e8_hf.png'))

sp6e8.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp7e8.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp9e8.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp10e8.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp13e8.specfit(fittype='gaussian', guesses=[1, 58, 1])
mlpi6e8 = sp6e8.specfit.parinfo
mlpi7e8 = sp7e8.specfit.parinfo
mlpi9e8 = sp9e8.specfit.parinfo
mlpi10e8 = sp9e8.specfit.parinfo
mlpi13e8 = sp9e8.specfit.parinfo
print("gaussian 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp6e8.specfit.parinfo))), file=outf)
print("gaussian 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp7e8.specfit.parinfo))), file=outf)
print("gaussian 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp9e8.specfit.parinfo))), file=outf)
print("gaussian 10-10 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp10e8.specfit.parinfo))), file=outf)
print("gaussian 13-13 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp13e8.specfit.parinfo))), file=outf)
"""
gaussian 6-6 peak=0.278613851186 +/- 0.00324112286397, fwhm=12.3888928639 +/- 0.166414956455, integral=19.3304597881 +/- 0.430595034151, center=59.0587207674 +/- 0.0706699251669
gaussian 7-7 peak=0.19046626119 +/- 0.00375662817502, fwhm=11.1985631899 +/- 0.255043061619, integral=10.7973493377 +/- 0.407787818476, center=59.7258654571 +/- 0.108306812353
gaussian 9-9 peak=0.0829717146021 +/- 0.00468072809892, fwhm=11.1981635835 +/- 0.729471923465, integral=4.70325110054 +/- 0.508070309565, center=59.6499937385 +/- 0.309778206795
"""

# e2w: absorption
sp6e2w,spK6e2w,_,_ = goddi_nh3_fits.load_spectrum(6, object='w51e2w', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp6e2w.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp7e2w,spK7e2w,_,_ = goddi_nh3_fits.load_spectrum(7, object='w51e2w', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp7e2w.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp9e2w,spK9e2w,_,_ = goddi_nh3_fits.load_spectrum(9, object='w51e2w', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp9e2w.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp10e2w,spK10e2w,_,_ = goddi_nh3_fits.load_spectrum(10, object='w51e2w', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp10e2w.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp13e2w,spK13e2w,_,_ = goddi_nh3_fits.load_spectrum(13, object='w51e2w', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp13e2w.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)

sp6e2w.plotter()
sp7e2w.plotter()
sp9e2w.plotter()
limits = [(50, 65), (25,30), (-10, 0), (0.1, 6), (30, 35), (-10, 0), (0.1, 6)]
limited = [(True,True)]*7
tied = ['']*7
tied[6] = 'p[3]'
tied[5] = 'p[2]'
sp6e2w.specfit(fittype='hfonly', guesses=[58, 26.9, -0.1, 2.0, 31.4, -0.1, 2.0], fixed=[F,F,F,F,F,F,F], limited=limited, limits=limits)
sp7e2w.specfit(fittype='hfonly', guesses=[58, 27.3, -0.1, 2.0, 31.2, -0.1, 2.0], fixed=[F,F,F,F,F,F,F], limited=limited, limits=limits)
sp9e2w.specfit(fittype='hfonly', guesses=[58, 27.0, -0.1, 2.0, 30.1, -0.1, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp6e2w.specfit.plot_fit(show_hyperfine_components=True)
sp7e2w.specfit.plot_fit(show_hyperfine_components=True)
sp9e2w.specfit.plot_fit(show_hyperfine_components=True)
hfpi6e2w = sp6e2w.specfit.parinfo
hfpi7e2w = sp7e2w.specfit.parinfo
hfpi9e2w = sp9e2w.specfit.parinfo
sp6e2w.plotter.savefig(os.path.join(savepath,'sp6e2w_hf.png'))
sp7e2w.plotter.savefig(os.path.join(savepath,'sp7e2w_hf.png'))
sp9e2w.plotter.savefig(os.path.join(savepath,'sp9e2w_hf.png'))

print("e2w", file=outf)
print("hf 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp6e2w.specfit.parinfo))), file=outf)
print("hf 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp7e2w.specfit.parinfo))), file=outf)
print("hf 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp9e2w.specfit.parinfo))), file=outf)
sp6e2w.specfit(fittype='gaussian', guesses=[-1, 58, 1])
sp7e2w.specfit(fittype='gaussian', guesses=[-1, 58, 1])
sp9e2w.specfit(fittype='gaussian', guesses=[-1, 58, 1])
sp10e2w.specfit(fittype='gaussian', guesses=[-1, 58, 1])
sp13e2w.specfit(fittype='gaussian', guesses=[-1, 58, 1])
mlpi6e2w = sp6e2w.specfit.parinfo
mlpi7e2w = sp7e2w.specfit.parinfo
mlpi9e2w = sp9e2w.specfit.parinfo
mlpi10e2w = sp9e2w.specfit.parinfo
mlpi13e2w = sp9e2w.specfit.parinfo
print("gaussian 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp6e2w.specfit.parinfo))), file=outf)
print("gaussian 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp7e2w.specfit.parinfo))), file=outf)
print("gaussian 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp9e2w.specfit.parinfo))), file=outf)
print("gaussian 10-10 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp10e2w.specfit.parinfo))), file=outf)
print("gaussian 13-13 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp13e2w.specfit.parinfo))), file=outf)

# e2e aka e2east
sp6e2e,spK6e2e,_,_ = goddi_nh3_fits.load_spectrum(6, object='w51e2e', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp6e2e.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
spK6e2e.specfit.Registry.add_fitter('sixsix',sixsix_movinghf_fitter(),7)
sp7e2e,spK7e2e,_,_ = goddi_nh3_fits.load_spectrum(7, object='w51e2e', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp7e2e.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp9e2e,spK9e2e,_,_ = goddi_nh3_fits.load_spectrum(9, object='w51e2e', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp9e2e.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp10e2e,spK10e2e,_,_ = goddi_nh3_fits.load_spectrum(10, object='w51e2e', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp10e2e.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp13e2e,spK13e2e,_,_ = goddi_nh3_fits.load_spectrum(13, object='w51e2e', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp13e2e.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)

sp6e2e.plotter()
spK6e2e.plotter()
sp7e2e.plotter()
sp9e2e.plotter()
# sp6e2e might be optically thick (in hyperfines, in emission!)
sp6e2e.specfit(fittype='hfonly', guesses=[58, 26.9, 0.02, 2.0, 31.4, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied)
spK6e2e.specfit(fittype='sixsix', guesses=[58, 100, 26.9, 31.4, 2.0, 1200, 0], fixed=[F,F,T,T,F,T,T])
spK6e2e.specfit.plot_fit(show_hyperfine_components=True)
sp7e2e.specfit(fittype='hfonly', guesses=[58, 27.3, 0.02, 2.0, 31.2, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied)
sp9e2e.specfit(fittype='hfonly', guesses=[58, 27.0, 0.02, 2.0, 30.1, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied)
sp6e2e.specfit.plot_fit(show_hyperfine_components=True)
sp7e2e.specfit.plot_fit(show_hyperfine_components=True)
sp9e2e.specfit.plot_fit(show_hyperfine_components=True)
hfpi6e2e = sp6e2e.specfit.parinfo
hfpi7e2e = sp7e2e.specfit.parinfo
hfpi9e2e = sp9e2e.specfit.parinfo
sp6e2e.plotter.savefig(os.path.join(savepath,'sp6e2e_hf.png'))
spK6e2e.plotter.savefig(os.path.join(savepath,'spK6e2e_hf.png'))
sp7e2e.plotter.savefig(os.path.join(savepath,'sp7e2e_hf.png'))
sp9e2e.plotter.savefig(os.path.join(savepath,'sp9e2e_hf.png'))

print("e2e", file=outf)
print("hf 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp6e2e.specfit.parinfo))), file=outf)
print("hf 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp7e2e.specfit.parinfo))), file=outf)
print("hf 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp9e2e.specfit.parinfo))), file=outf)
sp6e2e.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp7e2e.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp9e2e.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp10e2e.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp13e2e.specfit(fittype='gaussian', guesses=[1, 58, 1])
mlpi6e2e = sp6e2e.specfit.parinfo
mlpi7e2e = sp7e2e.specfit.parinfo
mlpi9e2e = sp9e2e.specfit.parinfo
mlpi10e2e = sp9e2e.specfit.parinfo
mlpi13e2e = sp9e2e.specfit.parinfo
print("gaussian 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp6e2e.specfit.parinfo))), file=outf)
print("gaussian 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp7e2e.specfit.parinfo))), file=outf)
print("gaussian 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp9e2e.specfit.parinfo))), file=outf)
print("gaussian 10-10 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp10e2e.specfit.parinfo))), file=outf)
print("gaussian 13-13 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp13e2e.specfit.parinfo))), file=outf)

# e2nw, e2northwest
sp6e2nw,spK6e2nw,_,_ = goddi_nh3_fits.load_spectrum(6, object='w51e2nw', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp6e2nw.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp7e2nw,spK7e2nw,_,_ = goddi_nh3_fits.load_spectrum(7, object='w51e2nw', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp7e2nw.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp9e2nw,spK9e2nw,_,_ = goddi_nh3_fits.load_spectrum(9, object='w51e2nw', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp9e2nw.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp10e2nw,spK10e2nw,_,_ = goddi_nh3_fits.load_spectrum(10, object='w51e2nw', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp10e2nw.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp13e2nw,spK13e2nw,_,_ = goddi_nh3_fits.load_spectrum(13, object='w51e2nw', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp13e2nw.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)

sp6e2nw.plotter()
sp7e2nw.plotter()
sp9e2nw.plotter()
sp6e2nw.specfit(fittype='hfonly', guesses=[58, 26.9, 0.02, 2.0, 31.4, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied)
sp7e2nw.specfit(fittype='hfonly', guesses=[58, 27.3, 0.02, 2.0, 31.2, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied)
limits = [(50, 65), (25,30), (0, 1), (0.1, 6), (30, 35), (0, 1), (0.1, 6)]
sp9e2nw.specfit(fittype='hfonly', guesses=[58, 27.0, 0.02, 2.0, 30.1, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], tied=tied, limited=limited, limits=limits)
sp6e2nw.specfit.plot_fit(show_hyperfine_components=True)
sp7e2nw.specfit.plot_fit(show_hyperfine_components=True)
sp9e2nw.specfit.plot_fit(show_hyperfine_components=True)
hfpi6e2nw = sp6e2nw.specfit.parinfo
hfpi7e2nw = sp7e2nw.specfit.parinfo
hfpi9e2nw = sp9e2nw.specfit.parinfo
sp6e2nw.plotter.savefig(os.path.join(savepath,'sp6e2nw_hf.png'))
sp7e2nw.plotter.savefig(os.path.join(savepath,'sp7e2nw_hf.png'))
sp9e2nw.plotter.savefig(os.path.join(savepath,'sp9e2nw_hf.png'))

print("e2nw", file=outf)
print("hf 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp6e2nw.specfit.parinfo))), file=outf)
print("hf 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp7e2nw.specfit.parinfo))), file=outf)
print("hf 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp9e2nw.specfit.parinfo))), file=outf)
sp6e2nw.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp7e2nw.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp9e2nw.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp10e2nw.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp13e2nw.specfit(fittype='gaussian', guesses=[1, 58, 1])
mlpi6e2nw = sp6e2nw.specfit.parinfo
mlpi7e2nw = sp7e2nw.specfit.parinfo
mlpi9e2nw = sp9e2nw.specfit.parinfo
mlpi10e2nw = sp9e2nw.specfit.parinfo
mlpi13e2nw = sp9e2nw.specfit.parinfo
print("gaussian 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp6e2nw.specfit.parinfo))), file=outf)
print("gaussian 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp7e2nw.specfit.parinfo))), file=outf)
print("gaussian 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp9e2nw.specfit.parinfo))), file=outf)
print("gaussian 10-10 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp10e2nw.specfit.parinfo))), file=outf)
print("gaussian 13-13 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp13e2nw.specfit.parinfo))), file=outf)

# e2clump
sp6e2clump,spK6e2clump,_,_ = goddi_nh3_fits.load_spectrum(6, object='w51e2-clump', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp6e2clump.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp7e2clump,spK7e2clump,_,_ = goddi_nh3_fits.load_spectrum(7, object='w51e2-clump', headerfile='/Users/adam/work/w51/goddi/W51-25GHzcont.map.image.fits')
sp7e2clump.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp9e2clump,spK9e2clump,_,_ = goddi_nh3_fits.load_spectrum(9, object='w51e2-clump', headerfile='/Users/adam/work/w51/goddi/W51-27GHzcont.map.image.fits')
sp9e2clump.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp10e2clump,spK10e2clump,_,_ = goddi_nh3_fits.load_spectrum(10, object='w51e2-clump', headerfile='/Users/adam/work/w51/goddi/W51-29GHzcont.map.image.fits')
sp10e2clump.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)
sp13e2clump,spK13e2clump,_,_ = goddi_nh3_fits.load_spectrum(13, object='w51e2-clump', headerfile='/Users/adam/work/w51/goddi/W51-33GHzcont.map.image.fits')
sp13e2clump.specfit.Registry.add_fitter('hfonly',hfonly_fitter(),7)

sp6e2clump.plotter()
sp7e2clump.plotter()
sp9e2clump.plotter()
limits = [(50, 65), (25,30), (0, 1), (0.1, 6), (30, 35), (0, 1), (0.1, 6)]
limited = [(True,True)]*7
sp6e2clump.specfit(fittype='hfonly', guesses=[58, 26.9, 0.02, 2.0, 31.4, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp7e2clump.specfit(fittype='hfonly', guesses=[58, 27.3, 0.02, 2.0, 31.2, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
sp9e2clump.specfit(fittype='hfonly', guesses=[58, 27.0, 0.02, 2.0, 30.1, 0.02, 2.0], fixed=[F,T,F,F,T,F,F], limited=limited, limits=limits, tied=tied)
hfpi6e2clump = sp6e2clump.specfit.parinfo
hfpi7e2clump = sp7e2clump.specfit.parinfo
hfpi9e2clump = sp9e2clump.specfit.parinfo
sp6e2clump.specfit.plot_fit(show_hyperfine_components=True)
sp7e2clump.specfit.plot_fit(show_hyperfine_components=True)
sp9e2clump.specfit.plot_fit(show_hyperfine_components=True)
sp6e2clump.plotter.savefig(os.path.join(savepath,'sp6e2clump_hf.png'))
sp7e2clump.plotter.savefig(os.path.join(savepath,'sp7e2clump_hf.png'))
sp9e2clump.plotter.savefig(os.path.join(savepath,'sp9e2clump_hf.png'))

print("e2clump", file=outf)
print("hf 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp6e2clump.specfit.parinfo))), file=outf)
print("hf 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp7e2clump.specfit.parinfo))), file=outf)
print("hf 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals(sp9e2clump.specfit.parinfo))), file=outf)
sp6e2clump.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp7e2clump.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp9e2clump.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp10e2clump.specfit(fittype='gaussian', guesses=[1, 58, 1])
sp13e2clump.specfit(fittype='gaussian', guesses=[1, 58, 1])
mlpi6e2clump = sp6e2clump.specfit.parinfo
mlpi7e2clump = sp7e2clump.specfit.parinfo
mlpi9e2clump = sp9e2clump.specfit.parinfo
mlpi10e2clump = sp9e2clump.specfit.parinfo
mlpi13e2clump = sp9e2clump.specfit.parinfo
print("gaussian 6-6 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp6e2clump.specfit.parinfo))), file=outf)
print("gaussian 7-7 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp7e2clump.specfit.parinfo))), file=outf)
print("gaussian 9-9 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp9e2clump.specfit.parinfo))), file=outf)
print("gaussian 10-10 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp10e2clump.specfit.parinfo))), file=outf)
print("gaussian 13-13 peak={0} +/- {1}, fwhm={2} +/- {3}, integral={4} +/- {5}, center={6} +/- {7}".format(*(tbl_vals_gaussian(sp13e2clump.specfit.parinfo))), file=outf)

outf.close()

def full_integral(pi_hf, pi_ml):
    """
    Combine the hyperfine and main line integral
    """
    if pi_hf is not None:
        i1,ei1 = integral_hf(pi_hf)
    else:
        i1,ei1 = 0,0

    i2,ei2 = integral_ml(pi_ml)

    return (i1+i2), (ei1**2+ei2**2)**0.5

def Nu_thin_hightex(flux, line_strength, freq, fillingfactor=1.0, tau=None):
    """
    Optically-thin-ish approximation for the column density of the upper state
    of a given line assuming T_ex >> T_bg and T_ex >> h nu
    """
    assert flux.unit.is_equivalent(u.K*u.km/u.s)
    assert line_strength.unit.is_equivalent(u.esu*u.cm)
    k = constants.k_B
    term1 = (3*k/(8*np.pi**2 * freq * line_strength**2))
    term5 = flux.to(u.K*u.km/u.s) / fillingfactor
    term6 = 1 if tau is None else tau/(1-np.exp(-tau))
    return (term1*term5*term6).to(u.cm**-2)

flux_e8 = {'66': (full_integral(hfpi6e8, mlpi6e8)*u.Jy).to(u.K, u.brightness_temperature(areas['e8'], freqs['66']))*u.km/u.s,
           '77': (full_integral(hfpi7e8, mlpi7e8)*u.Jy).to(u.K, u.brightness_temperature(areas['e8'], freqs['77']))*u.km/u.s,
           '99': (full_integral(hfpi9e8, mlpi9e8)*u.Jy).to(u.K, u.brightness_temperature(areas['e8'], freqs['99']))*u.km/u.s,
           '1010': (full_integral(None, mlpi10e8)*u.Jy).to(u.K, u.brightness_temperature(areas['e8'], freqs['1010']))*u.km/u.s,
           '1313': (full_integral(None, mlpi13e8)*u.Jy).to(u.K, u.brightness_temperature(areas['e8'], freqs['1313']))*u.km/u.s,
          }
flux_e2e = {'66': (full_integral(hfpi6e2e, mlpi6e2e)*u.Jy).to(u.K, u.brightness_temperature(areas['e2e'], freqs['66']))*u.km/u.s,
            '77': (full_integral(hfpi7e2e, mlpi7e2e)*u.Jy).to(u.K, u.brightness_temperature(areas['e2e'], freqs['77']))*u.km/u.s,
            '99': (full_integral(hfpi9e2e, mlpi9e2e)*u.Jy).to(u.K, u.brightness_temperature(areas['e2e'], freqs['99']))*u.km/u.s,
            '1010': (full_integral(None, mlpi10e2e)*u.Jy).to(u.K, u.brightness_temperature(areas['e2e'], freqs['1010']))*u.km/u.s,
            '1313': (full_integral(None, mlpi13e2e)*u.Jy).to(u.K, u.brightness_temperature(areas['e2e'], freqs['1313']))*u.km/u.s,
          }
taue2e = {'66': 100,
            '77': 77,
            '99': 43,
            '1010': None,
            '1313': None,
           }
taue8 = {'66':27,
         '77':33,
         '99':None,
         '1010': None,
         '1313': None,
           }
# if you want to exclude 1010 from the fit, uncomment this
#freqs.pop('1010')

Nu_e8 = {line: Nu_thin_hightex(flux_e8[line], linestrengths[line], freqs[line]) for line in freqs}
Nu_e2e = {line: Nu_thin_hightex(flux_e2e[line], linestrengths[line], freqs[line]) for line in freqs}
Nu_e8_taucorr = {line: Nu_thin_hightex(flux_e8[line], linestrengths[line], freqs[line], tau=taue8[line]) for line in freqs}
Nu_e2e_taucorr = {line: Nu_thin_hightex(flux_e2e[line], linestrengths[line], freqs[line], tau=taue2e[line]) for line in freqs}

flux_e2nw = {'66': (full_integral(hfpi6e2nw, mlpi6e2nw)*u.Jy).to(u.K, u.brightness_temperature(areas['e2nw'], freqs['66']))*u.km/u.s,
            '77': (full_integral(hfpi7e2nw, mlpi7e2nw)*u.Jy).to(u.K, u.brightness_temperature(areas['e2nw'], freqs['77']))*u.km/u.s,
            '99': (full_integral(hfpi9e2nw, mlpi9e2nw)*u.Jy).to(u.K, u.brightness_temperature(areas['e2nw'], freqs['99']))*u.km/u.s,
#            '1010': (full_integral(None, mlpi10e2nw)*u.Jy).to(u.K, u.brightness_temperature(areas['e2nw'], freqs['1010']))*u.km/u.s,
#            '1313': (full_integral(None, mlpi13e2nw)*u.Jy).to(u.K, u.brightness_temperature(areas['e2nw'], freqs['1313']))*u.km/u.s,
          }
taue2nw = {'66': None,
            '77': None,
            '99': None,
            '1010': None,
            '1313': None,
           }
Nu_e2nw = {line: Nu_thin_hightex(flux_e2nw[line], linestrengths[line], freqs[line]) for line in freqs if line in flux_e2nw}

flux_e2clump = {'66': (full_integral(hfpi6e2clump, mlpi6e2clump)*u.Jy).to(u.K, u.brightness_temperature(areas['e2clump'], freqs['66']))*u.km/u.s,
            '77': (full_integral(hfpi7e2clump, mlpi7e2clump)*u.Jy).to(u.K, u.brightness_temperature(areas['e2clump'], freqs['77']))*u.km/u.s,
            '99': (full_integral(hfpi9e2clump, mlpi9e2clump)*u.Jy).to(u.K, u.brightness_temperature(areas['e2clump'], freqs['99']))*u.km/u.s,
#            '1010': (full_integral(None, mlpi10e2clump)*u.Jy).to(u.K, u.brightness_temperature(areas['e2clump'], freqs['1010']))*u.km/u.s,
#            '1313': (full_integral(None, mlpi13e2clump)*u.Jy).to(u.K, u.brightness_temperature(areas['e2clump'], freqs['1313']))*u.km/u.s,
          }
taue2clump = {'66': None,
            '77': None,
            '99': None,
            '1010': None,
            '1313': None,
           }
Nu_e2clump = {line: Nu_thin_hightex(flux_e2clump[line], linestrengths[line], freqs[line]) for line in freqs if line in flux_e2clump}

# http://spec.jpl.nasa.gov/ftp/pub/catalog/doc/d017002.pdf
B0=(298117.0*u.MHz)
C0=(186726.0*u.MHz)

from astropy import modeling

def fit_tex(eupper, nupperoverg):
    """
    Fit the Boltzmann diagram
    """
    model = modeling.models.Linear1D()
    fitter = modeling.fitting.LevMarLSQFitter()
    result = fitter(model, eupper, np.log(nupperoverg))
    tex = -1./result.slope*u.K

    #calculating partition function Q_rot; use eqn 15.48 tools of radio astronomy
    # = sqrt( (pi (kT)^3) / h^3 B^2 C )
    Q_rot= np.sqrt(np.pi * (constants.k_B*tex)**3 / (constants.h**3 * B0**2 * C0)).decompose()

    Ntot = np.exp(result.intercept + np.log(Q_rot)) * u.cm**-2

    print(("Tex={0}, Ntot={1}, Q_rot={2}".format(tex, Ntot, Q_rot)))

    return Ntot, tex, result.slope, result.intercept

import pylab as pl
pl.matplotlib.rc_file('pubfiguresrc')

eups = np.array([eupper[line] for line in freqs])

fig = pl.figure(1)
fig.clf()
ax = pl.gca()
print("e8")
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs],
            y=[Nu_e8[line][0].value*degeneracy[line] for line in freqs],
            yerr=[Nu_e8[line][1].value*degeneracy[line] for line in freqs],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs], [Nu_e8[line][0].value*degeneracy[line] for line in freqs])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
print(("Mass lower limit: {0}".format((Ntot/XNH3 * areas['e8']*distance**2 * 2.8*u.Da).to(u.M_sun, u.dimensionless_angles()))))
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs],
            y=[Nu_e8_taucorr[line][0].value*degeneracy[line] for line in freqs],
            yerr=[Nu_e8_taucorr[line][1].value*degeneracy[line] for line in freqs],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs],
                                [Nu_e8_taucorr[line][0].value*degeneracy[line] for line in freqs])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
ax.set_yscale('log')
#ax.set_xscale('log')
ax.set_ylabel("$N_u / g$ [cm$^{-2}$]")
ax.set_xlabel("$E_u$ [K]")
ax.set_title('e8')
pl.legend(loc='best')
pl.savefig(os.path.join(savepath, 'e8_rotational_diagram_fit.png'))

print("e2e")
fig = pl.figure(2)
fig.clf()
ax = pl.gca()
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs],
            y=[Nu_e2e[line][0].value*degeneracy[line] for line in freqs],
            yerr=[Nu_e2e[line][1].value*degeneracy[line] for line in freqs],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs], [Nu_e2e[line][0].value*degeneracy[line] for line in freqs])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
print(("Mass lower limit: {0}".format((Ntot/XNH3 * areas['e2e']*distance**2 * 2.8*u.Da).to(u.M_sun, u.dimensionless_angles()))))
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs],
            y=[Nu_e2e_taucorr[line][0].value*degeneracy[line] for line in freqs],
            yerr=[Nu_e2e_taucorr[line][1].value*degeneracy[line] for line in freqs],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs],
                                [Nu_e2e_taucorr[line][0].value*degeneracy[line] for line in freqs])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
ax.set_yscale('log')
#ax.set_xscale('log')
ax.set_ylabel("$N_u / g$ [cm$^{-2}$]")
ax.set_xlabel("$E_u$ [K]")
ax.set_title('e2e')
pl.legend(loc='best')
pl.savefig(os.path.join(savepath, 'e2e_rotational_diagram_fit.png'))

print("e2nw")
fig = pl.figure(3)
fig.clf()
ax = pl.gca()
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs if line in Nu_e2nw],
            y=[Nu_e2nw[line][0].value*degeneracy[line] for line in freqs if line in Nu_e2nw],
            yerr=[Nu_e2nw[line][1].value*degeneracy[line] for line in freqs if line in Nu_e2nw],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs if line in Nu_e2nw], [Nu_e2nw[line][0].value*degeneracy[line] for line in freqs if line in Nu_e2nw])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
print(("Mass lower limit: {0}".format((Ntot/XNH3 * areas['e2nw']*distance**2 * 2.8*u.Da).to(u.M_sun, u.dimensionless_angles()))))
ax.set_yscale('log')
#ax.set_xscale('log')
ax.set_ylabel("$N_u / g$ [cm$^{-2}$]")
ax.set_xlabel("$E_u$ [K]")
ax.set_title('e2nw')
pl.legend(loc='best')
pl.savefig(os.path.join(savepath, 'e2nw_rotational_diagram_fit.png'))

print("e2clump")
fig = pl.figure(4)
fig.clf()
ax = pl.gca()
a,b,c = ax.errorbar(x=[eupper[line] for line in freqs if line in Nu_e2clump],
            y=[Nu_e2clump[line][0].value*degeneracy[line] for line in freqs if line in Nu_e2clump],
            yerr=[Nu_e2clump[line][1].value*degeneracy[line] for line in freqs if line in Nu_e2clump],
            marker='s',
            linestyle='none')
Ntot,tex,slope,intcpt = fit_tex([eupper[line] for line in freqs if line in Nu_e2clump], [Nu_e2clump[line][0].value*degeneracy[line] for line in freqs if line in Nu_e2clump])
ax.plot(eups, np.exp(eups*slope + intcpt), color=a.get_color(), label='$T_{{ex}}={0:0.1f}$\n$N(\\mathrm{{NH}}_3)={1:0.1e}$ cm$^{{-2}}$'.format(tex, Ntot.value))
print(("Mass lower limit: {0}".format((Ntot/XNH3 * areas['e2clump']*distance**2 * 2.8*u.Da).to(u.M_sun, u.dimensionless_angles()))))
ax.set_yscale('log')
#ax.set_xscale('log')
ax.set_ylabel("$N_u / g$ [cm$^{-2}$]")
ax.set_xlabel("$E_u$ [K]")
ax.set_title('e2clump')
pl.legend(loc='best')
pl.savefig(os.path.join(savepath, 'e2clump_rotational_diagram_fit.png'))

pl.draw()
pl.show()