import cv2
import numpy as np
import math
import RPi.GPIO as GPIO

# Motors
GPIO.setwarnings(False)
GPIO.setmode(GPIO.BOARD)

motor_1_pin = 40 #FIND PWM PIN
motor_1_1 = 29
motor_1_2 = 31

motor_2_pin = 38 #ABOVE
motor_2_1 = 16
motor_2_2 = 18

default_speed = 90

# Image
rho = 1
theta = np.pi / 180
min_threshold = 10
minLineLength = 5
maxLineGap = 20

video = cv2.VideoCapture(0)

video.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
video.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)

def convert_to_hsv(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    # cv2.imshow("HSV", hsv)

    return hsv

def detect_edges(frame):
    lower_limit = np.array([90, 120, 0], dtype="uint8")  # lower limit of blue
    upper_limit = np.array([150, 255, 255], dtype="uint8")  # upper limit of blue
    mask = cv2.inRange(hsv, lower_limit, upper_limit)  # filter for colour limits

    edges = cv2.Canny(mask, 50, 100)
    # cv2.imshow("edges", edges)

    return edges

def select_roi(edges):
    height, width = edges.shape
    mask = np.zeros_like(edges)  # empty matrix with same dimensions as frame

    # focus on lower half of screen
    # specify coordinates of 4 points(lower left, upper left, upper right, lower right)
    polygon = np.array([[
        (0, height),
        (0, height/2),
        (width, height/2),
        (width, height),
    ]], np.int32)

    cv2.fillPoly(mask, polygon, 255)  # fill polygon with blue
    cropped_edges = cv2.bitwise_and(edges, mask)
    # cv2.imshow("roi", cropped_edges)

    return cropped_edges

def detect_line_segments(cropped_edges):
    line_segments = cv2.HoughLinesP(cropped_edges, rho, theta, min_threshold,
                                   np.array([]), minLineLength, maxLineGap)

    return line_segments

def make_points(frame, line):
    height, width, _ = frame.shape
    slope, intercept = line
    y1 = height  # bottom of the frame
    y2 = int(y1 / 2)  # make points from middle of the frame down

    if slope == 0:
        slope = 0.1

    x1 = int((y1 - intercept) / slope)
    x2 = int((y2 - intercept) / slope)

    return [[x1, y1, x2, y2]]

def display_lines(frame, lines, line_color=(0, 255, 0), line_width=6):  # line color (B,G,R)
    line_image = np.zeros_like(frame)
    if lines is not None:
        for line in lines:
            for x1, y1, x2, y2 in line:
                cv2.line(line_image, (x1, y1), (x2, y2), line_color, line_width)
    line_image = cv2.addWeighted(frame, 0.8, line_image, 1, 1)

    return line_image

def average_slope_intercept(frame, line_segments):
    lane_lines = []
    if line_segments is None:
        print("no line segment detected")
        return lane_lines
    height, width,_ = frame.shape
    left_fit = []
    right_fit = []
    boundary = 1/3
    left_region_boundary = width * (1 - boundary)
    right_region_boundary = width * boundary
    for line_segment in line_segments:
        for x1, y1, x2, y2 in line_segment:
            if x1 == x2:
                print("skipping vertical lines (slope = infinity)")
                continue
            fit = np.polyfit((x1, x2), (y1, y2), 1)
            slope = (y2 - y1) / (x2 - x1)
            intercept = y1 - (slope * x1)
            if slope < 0:
                if x1 < left_region_boundary and x2 < left_region_boundary:
                    left_fit.append((slope, intercept))
            else:
                if x1 > right_region_boundary and x2 > right_region_boundary:
                    right_fit.append((slope, intercept))
    left_fit_average = np.average(left_fit, axis=0)
    if len(left_fit) > 0:
        lane_lines.append(make_points(frame, left_fit_average))
    right_fit_average = np.average(right_fit, axis=0)
    if len(right_fit) > 0:
        lane_lines.append(make_points(frame, right_fit_average))
    # lane_lines is a 2-D array consisting the coordinates of the right and left lane lines
    # for example: lane_lines = [[x1,y1,x2,y2],[x1,y1,x2,y2]]
    # where the left array is for left lane and the right array is for right lane
    # all coordinate points are in pixels
    return lane_lines

def get_steering_angle(frame, lane_lines):
    height, width, _ = frame.shape
    x_offset, y_offset = 0, 0

    if len(lane_lines) == 2: # if two lane lines are detected
        _, _, left_x2, _ = lane_lines[0][0] # extract left x2 from lane_lines array
        _, _, right_x2, _ = lane_lines[1][0] # extract right x2 from lane_lines array
        mid = int(width / 2)
        x_offset = (left_x2 + right_x2) / 2 - mid
        y_offset = int(height / 2)

    elif len(lane_lines) == 1: # if only one line is detected
        x1, _, x2, _ = lane_lines[0][0]
        x_offset = x2 - x1
        y_offset = int(height / 2)

    elif len(lane_lines) == 0: # if no line is detected
        x_offset = 0
        y_offset = int(height / 2)

    angle_to_mid_radian = math.atan(x_offset / y_offset)
    angle_to_mid_deg = int(angle_to_mid_radian * 180.0 / math.pi)
    steering_angle = angle_to_mid_deg + 90

    return steering_angle

def display_heading_line(frame, steering_angle, line_color=(0, 0, 255), line_width=5):
    heading_image = np.zeros_like(frame)
    height, width, _ = frame.shape
    steering_angle_radian = steering_angle / 180.0 * math.pi
    x1 = int(width / 2)
    y1 = height
    x2 = int(x1 - height / 2 / math.tan(steering_angle_radian))
    y2 = int(height / 2)
    cv2.line(heading_image, (x1, y1), (x2, y2), line_color, line_width)
    heading_image = cv2.addWeighted(frame, 0.8, heading_image, 1, 1)
    cv2.imshow("heading", heading_image)

    return heading_image

if __name__ == '__main__':
    # init
    GPIO.setup(motor_1_pin, GPIO.OUT)
    GPIO.setup(motor_2_pin, GPIO.OUT)
    GPIO.setup(motor_1_1, GPIO.OUT)
    GPIO.setup(motor_1_2, GPIO.OUT)
    GPIO.setup(motor_2_1, GPIO.OUT)
    GPIO.setup(motor_2_2, GPIO.OUT)

    GPIO.output(motor_1_1, True)
    GPIO.output(motor_1_2, False)
    GPIO.output(motor_2_1, False)
    GPIO.output(motor_2_2, True)

    motor_1 = GPIO.PWM(motor_1_pin, 1000)
    motor_1.stop()

    motor_2 = GPIO.PWM(motor_2_pin, 1000)
    motor_2.stop()

    while True:
        # Lane detection
        ret, frame = video.read()
        # frame = cv2.flip(frame,-1) # used to flip the image vertically
        hsv = convert_to_hsv(frame)  # convert to hsv
        edges = detect_edges(hsv)  # detect colour
        roi_edges = select_roi(edges)  # returns only region of interest
        line_segments = detect_line_segments(roi_edges)  # detects lines in roi
        lane_lines = average_slope_intercept(frame, line_segments)
        lane_lines_image = display_lines(frame, lane_lines)
        steering_angle = get_steering_angle(frame, lane_lines)
        heading_image = display_heading_line(lane_lines_image, steering_angle)
        # cv2.imshow('original', frame)

        print(steering_angle);
        # Motors
        motor_1.start(default_speed)
        motor_2.start(default_speed)
        #if steering_angle == 90:
        #    motor_1.start(default_speed)
        #    motor_2.start(default_speed)
        #elif steering_angle > 90:
        #    motor_1.start(steering_angle-90)
        #    motor_2.start(default_speed)
        #elif steering_angle < 90:
        #    motor_1.start(default_speed)
        #    motor_2.start(steering_angle)

        # Exit key
        key = cv2.waitKey(1)
        if key == 27:
            break

    video.release()
    cv2.destroyAllWindows()


# -*- coding:utf-8 -*-
import os
import time
import pymongo
import requests
from dotenv import load_dotenv
load_dotenv()

def get_collection():
    mongo_url = os.getenv('MONGO_URL')
    client = pymongo.MongoClient(mongo_url)
    collection = client.image_collection
    return collection

def main():
    collection = get_collection()
    images = collection.images

    while True:
        try:
            _cpu_serial_id = None
            file_list = []
            for meta in images.find():
                _id = meta['_id']
                _cpu_serial_id = meta['cpu_serial']
                _image_path = meta['path']

                file_list.append(('image', open(_image_path, 'rb')))

                # remove information (mongo, database)
                images.delete_one({'_id': _id})
                os.remove(_image_path)

            if len(file_list) == 0:
                continue

            url = os.getenv('TARGET_URL') + 'datacollections/image/'
            res = requests.post(
                url,
                data={'cpu_serial_id': _cpu_serial_id},
                files=file_list
            )
        except KeyboardInterrupt:
            break

if __name__ == '__main__':
    main()


# Schedule Announcer Mate

# Imports

import datetime

from google.oauth2 import service_account
from googleapiclient.discovery import build

from threading import Thread

import pyttsx3

import time

# Get Structured Date Function
def get_current_structured_date():
    current_datetime = datetime.datetime.now()
    current_day_of_the_month = current_datetime.strftime("%d")
    current_month = current_datetime.strftime("%B")
    current_year = current_datetime.strftime("%Y")
    current_date = f"{current_day_of_the_month} {current_month} {current_year}"
    return current_date

# Get Structured Time Function
def get_current_structured_time():
    current_datetime = datetime.datetime.now()
    current_hour = current_datetime.strftime("%H")
    current_minute = current_datetime.strftime("%M")
    current_time = f"{current_hour}:{current_minute}"
    return [current_time, current_hour]

# pyttsx3 Variables

saying_engine = pyttsx3.init()
saying_engine.setProperty("rate", 200)

speaker_voices = saying_engine.getProperty("voices")
saying_engine.setProperty("voice", speaker_voices[1].id)

# For Sam To Be Able To Speak
def speak(thing_to_say):
    saying_engine.say(thing_to_say)
    saying_engine.runAndWait()

# Google Sheets Variables

credentials = service_account.Credentials.from_service_account_file("../credentials.json",  scopes=['https://www.googleapis.com/auth/spreadsheets'])
service = build("sheets", "v4", credentials=credentials)

spreadsheet_id = "1R1ooZQYaSHPByec1VqjOjUUd1SFjH_YsCsgpeZGXssI"

# Getting Schedule Accordingly From Google Sheets
def get_schedule(schedule_date):
    sheet_name = schedule_date
    schedule_range = f"{sheet_name}!B3:D290"

    # Response Structure
    # response = {values: arrays_of(time, task, task_details)}

    response = service.spreadsheets().values().get(spreadsheetId = spreadsheet_id, range = schedule_range).execute()
    return response

# Dictionary For Announcements With Their Times

# Structure = {"time": {"task": "task details"}}
CURRENT_SCHEDULE = {}
SAYING_QUEUE = []

# Function For Populating Global Schedule Dictionary With A Proper Format

def populate_current_schedule(schedule_to_populate_with):

    global CURRENT_SCHEDULE

    schedule_values = schedule_to_populate_with["values"]

    current_hour = get_current_structured_time()[1]

    for schedule_value in schedule_values:

        length_of_schedule_value_array = len(schedule_value)

        if length_of_schedule_value_array != 1:
            task_time = schedule_value[0]
            task = schedule_value[1]

            task_hour = task_time.split(":")[0]

            if task_hour >= current_hour:

                if length_of_schedule_value_array == 3:
                    task_detail = schedule_value[2]
                    CURRENT_SCHEDULE[task_time] = f"{task}, {task_detail}"
                else:
                    CURRENT_SCHEDULE[task_time] = f"{task}"

# Thread For Properly Announcing
def announcing_stuff():

    global SAYING_QUEUE

    while True:
        for thing_to_say in SAYING_QUEUE:
            speak(thing_to_say)

announcing_stuff_thread = Thread(target=announcing_stuff)

# Thread For Checking And Updating Date | Accordingly, Configuring The Dictionary

# To Check:
    # Date Changes
    # New Changes In Sheet

def keeping_schedule_up_to_date():

    previous_date = get_current_structured_date()
    previous_schedule = get_schedule(previous_date)

    populate_current_schedule(previous_schedule)

    while True:

        new_date = get_current_structured_date()
        new_schedule = get_schedule(new_date)

        if new_date != previous_date:

            populate_current_schedule(new_schedule)
            previous_date = new_date

        else:

            if new_schedule != previous_schedule:
                populate_current_schedule(new_schedule)
                previous_schedule = new_schedule

        time.sleep(2)

thread_for_keeping_schedule_up_to_date = Thread(target=keeping_schedule_up_to_date)

# Thread For Schedule Announcement For The Day | Going Through The Dictiornary, Checking For Each's Time

def pushing_schedule_to_announcement():

    global CURRENT_SCHEDULE, SAYING_QUEUE

    while True:

        current_time = get_current_structured_time()[0]

        try:

            for task_time, task_announcement in CURRENT_SCHEDULE.items():

                if current_time == task_time:

                    CURRENT_SCHEDULE.pop(task_time)

                    for iteration in range(3):
                        SAYING_QUEUE.append(task_announcement)

        except RuntimeError:

            error = "Current Schedule Was Changed When It Was Being Operated On"

pushing_schedule_to_announcement = Thread(target=pushing_schedule_to_announcement)

# Main With Everything Compiled

if __name__ == "__main__":
    announcing_stuff_thread.start()
    thread_for_keeping_schedule_up_to_date.start()
    pushing_schedule_to_announcement.start()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec  8 14:21:40 2020

@author: gilleschardon
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
import acosolo.cmf as cmf
from acosolo.sourcemodels import freefield2D
from acosolo.beamforming import bmap_unc
from acosolo.utils import grid3D

# CMF-NNLS

# Ref.: G. Chardon, J. Picheral, F. Ollivier, Theoretical analysis of the DAMAS algorithm and efficient implementation of the Covariance Matrix Fitting method for large-scale problems, Journal of Sound and Vibration,

mat = loadmat("damasdata2D.mat")
k = mat['k'][0,0]
Sigma = mat['Data']
Array = mat['Pmic']

# Normalization, for stability
Sigma = Sigma / np.max(Sigma)

# Source grid
Z = 4.6

g2D = lambda x : freefield2D(Array, x, Z, k)

Xgrid, dimgrid = grid3D([-2, -1, Z], [1, 0, Z], 0.02)
A = g2D(Xgrid)
bmap = bmap_unc(Sigma, A)

cmf_map = cmf.cmf_nnls_lh(A, Sigma, dr = True)

cmfDB = 10*np.log10(cmf_map)

def plotmap(pmap, name, dynrange):
    mapDB = 10*np.log10(pmap)
    m = np.max(mapDB)
    plt.figure()
    plt.imshow(mapDB, cmap='hot', vmax=m, vmin=m-dynrange, origin="lower")
    current_cmap = plt.cm.get_cmap()
    current_cmap.set_bad(color='black')
    plt.axis('image')
    plt.title(name)
    ax = plt.gca()
    ax.set_facecolor('black')

#%%

plotmap(np.reshape(cmf_map, [dimgrid[1], dimgrid[0]]), 'CMF-NNLS', 20)
plotmap(np.reshape(bmap, [dimgrid[1], dimgrid[0]]), 'Beamforming', 20)


import markups, config, Forms, create_bot
from datetime import datetime
from aiogram import types, Dispatcher
from create_bot import dp, bot, db
from aiogram.dispatcher import FSMContext

async def start(message : types.Message):
    print('/start from some user', message.from_user.username)
    user = db.get_user_info(message.from_user.id)
    if not user:
        try: print("User - {0.username}[{0.id}] started use the bot at {1}".format(message.from_user, datetime.now().strftime('%H:%M:%S:%f %d/%m/%Y')))
        except Exception: pass
        db.sqlite(f"INSERT INTO users (user_id, full_name, is_passed_test, level) VALUES ({message.from_user.id}, '{message.from_user.full_name}', 0, 'Junior')")
        cell = await create_bot.user_workspace.find(str(message.from_user.id), in_column=9)
        try:
            member = await bot.get_chat_member("@angli3i", message.from_user.id)
            isSub = 'Так' if member.is_chat_member() else 'Ні'
        except Exception: isSub = 'Невідомо'
        if cell is None:
            await create_bot.user_workspace.append_row([
                datetime.now().strftime('%d.%m.%Y - %H:%M'),
                None,
                message.from_user.full_name,
                message.from_user.username,
                None,
                'Так',
                0,
                isSub,
                message.from_user.id,
                'Junior',
            ])
        await message.answer_video_note(
            open('WelcomePageItems/doc_2023-01-03_18-37-37.mp4', 'rb')
        )
        if message.from_user.full_name: await message.answer(
            f'Вітаю, *{message.from_user.full_name}*!',
            parse_mode='Markdown'
        )

        user = db.get_user_info(message.from_user.id)

    await message.answer_photo(
        open('Menu/Frame 39.png', 'rb'),
        reply_markup=markups.MainMenu if (user and user[5] > 0) or db.sqlite(f"SELECT * FROM GameUsers WHERE user_id = {message.from_user.id}") else markups.Menu
    )

async def start_menu(callback : types.CallbackQuery):
    user=db.get_user_info(callback.from_user.id)
    await callback.message.edit_media(
        types.InputMediaPhoto(types.InputFile('Menu/Frame 39.png')),
        reply_markup=markups.MainMenu if (user and user[5] > 0) or db.sqlite(f"SELECT * FROM GameUsers WHERE user_id = {callback.from_user.id}") else markups.Menu
    )

async def angli3i_info(callback : types.CallbackQuery):
    if callback.data != 'menu_back': await callback.message.answer_video_note(open('WelcomePageItems/doc_2023-01-03_18-35-30.mp4', 'rb'))
    try: await callback.message.delete()
    except Exception: pass
    await callback.message.answer_photo(
        open('Menu/Frame 45.png', 'rb'),
        reply_markup=markups.angli3iMenu
    )
    db.sqlite(f"UPDATE Users SET step = step + 1 WHERE user_id = {callback.from_user.id}")
    step, = db.get_user_info(callback.from_user.id, 'step')
    cell = await create_bot.user_workspace.find(str(callback.from_user.id), in_column=9)
    if cell is not None:
        await create_bot.user_workspace.update_cell(cell.row, config.angli3i_steps.get(callback.data), step)

async def selected_game(callback : types.CallbackQuery):
    num_played_games, level = db.get_user_info(callback.from_user.id, 'num_played_games, level')
    game_code = callback.data.split('_')[-1]
    price = 1 if num_played_games < 1 else config.prices.get(level, 1)
    game = db.get_game(game_code)

    markup = types.InlineKeyboardMarkup(inline_keyboard=[
        [types.InlineKeyboardButton("ЗАРЕЄСТРУВАТИСЬ", pay=True)]
    ])
    prices = [
        types.LabeledPrice(label="Реєстрація на гру", amount=price*100)
    ]
    await callback.message.delete()
    await bot.send_invoice(
            callback.message.chat.id,
            title=game[0],
            description=game[1],
            provider_token=config.PAYMENTPROVIDERTOKEN,
            currency="UAH",
            prices=prices,
            start_parameter=f"reg_game_{game_code}",
            payload=f"reg_game_{game_code}",
            reply_markup=markup
        )

async def open_menu_angli3i(callback : types.CallbackQuery):
    data = callback.data.split('_')[-1]
    video = markups.buttons.get(data).get('video')
    photo = markups.buttons.get(data).get('photo')

    try: await callback.message.delete()
    except Exception: pass
    if video: await callback.message.answer_video_note(open(video,'rb'))
    await callback.message.answer_photo(
        "https://raw.githubusercontent.com/EvenEasy/EngEasyBot1/master/" + photo,
        reply_markup=markups.angli3iMenu if callback.data != 'menu_whoplays' else markups.who_plays_markup
    )

    db.sqlite(f"UPDATE Users SET step = step + 1 WHERE user_id = {callback.from_user.id}")
    step, = db.get_user_info(callback.from_user.id, 'step')
    cell = await create_bot.user_workspace.find(str(callback.from_user.id), in_column=9)
    if cell is not None:
        await create_bot.user_workspace.update_cell(cell.row, config.angli3i_steps.get(callback.data), step)

async def reg_welcome_page(callback : types.CallbackQuery):
    num_played_games, level = db.get_user_info(callback.from_user.id, "num_played_games, level")
    game = db.get_new_game(level)
    if db.sqlite(f"SELECT * FROM GameUsers WHERE user_id = {callback.from_user.id}"):
        await callback.message.edit_media(
            types.InputMediaPhoto(types.InputFile('Menu/photo_2023-02-03_19-34-13.jpg')),
            callback.message.reply_markup
        )
        try: await callback.answer()
        except Exception: pass
        return

    price = 1 # if num_played_games < 1 else config.prices.get(level, 1)

    if not game and num_played_games > 0:
        await callback.message.edit_media(
            types.InputMediaPhoto(types.InputFile('Menu/photo_2023-02-09_19-28-34.jpg')),
            reply_markup=callback.message.reply_markup
        )
        return

    try:
        if game and game[3] is not None:
            game_limit, = db.get_game(game[3], 'limit_players')
            nums_players = len(db.sqlite(f"SELECT user_id FROM GameUsers WHERE game_code = '{game[3]}'"))
            if game_limit <= nums_players:
                await callback.message.edit_media(
                    types.InputMediaPhoto(types.InputFile('Menu/photo_2023-02-06_20-59-16.jpg')),
                    callback.message.reply_markup
                )
                return
        prices = [
            types.LabeledPrice(label="Реєстрація на гру", amount=price*100)
        ]
        markup = types.InlineKeyboardMarkup(inline_keyboard=[
            [types.InlineKeyboardButton(f"ЗАРЕЄСТРУВАТИСЬ ({nums_players}/{game_limit})" if game else "ЗАРЕЄСТРУВАТИСЬ", pay=True)]
        ])

        if num_played_games == 0:
            markup.add(types.InlineKeyboardButton("УГОДА КОРИСТУВАЧА", url="https://telegra.ph/UGODA-KORISTUVACHA-01-16"))
            markup.add(types.InlineKeyboardButton("ПОЛіТИКА КОНФіДЕНЦіЙНОСТі", url="https://telegra.ph/POLІTIKA-KONFІDENCІJNOSTІ-01-16"))

        markup.add(types.InlineKeyboardButton('назад', callback_data='back_to_menu'))

        await callback.message.delete()
        await bot.send_invoice(
            callback.message.chat.id,
            title=game[0],
            description=game[1],
            provider_token=config.PAYMENTPROVIDERTOKEN,
            currency="UAH",
            prices=prices,
            start_parameter=f"reg_game_{game[3]}",
            payload=f"reg_game_{game[3]}",
            reply_markup=markup
        )

    except Exception as E:
        print(str(E))
        await callback.message.answer("Сталася якась помилка з платіжною системою")

async def reg_game(callback : types.CallbackQuery):
    game_code = callback.data.split('_')[2::]
    game_code, level = game_code
    num_played_games, = db.get_user_info(callback.from_user.id, 'num_played_games')
    game = db.get_game(game_code)
    price = 1 if num_played_games < 1 else config.prices.get(level, 1)

    try:
        prices = [
            types.LabeledPrice(label="Реєстрація на гру", amount=price*100)
        ]
        markup = types.InlineKeyboardMarkup(inline_keyboard=[
            [types.InlineKeyboardButton("ЗАРЕЄСТРУВАТИСЬ", pay=True)]
        ])
        if price == 1:
            markup.add(types.InlineKeyboardButton("УГОДА КОРИСТУВАЧА", url="https://telegra.ph/UGODA-KORISTUVACHA-01-16"))
            markup.add(types.InlineKeyboardButton("ПОЛіТИКА КОНФіДЕНЦіЙНОСТі", url="https://telegra.ph/POLІTIKA-KONFІDENCІJNOSTІ-01-16"))

        markup.add(types.InlineKeyboardButton('назад', callback_data='back_to_menu'))

        await callback.message.delete()
        await bot.send_invoice(
            callback.message.chat.id,
            title=game[0],
            description=game[1],
            provider_token=config.PAYMENTPROVIDERTOKEN,
            currency="UAH",
            prices=prices,
            start_parameter=f"reg_game_{game_code}",
            payload=f"reg_game_{game_code}",
            reply_markup=markup
        )

    except Exception as E:
        print(str(E))
        await callback.message.answer("Сталася якась помилка з платіжною системою")

async def time_to_game(message : types.Message):
    data_time = db.sqlite(f"""
        SELECT Games._date, Games.title, Games.description
        FROM Games
        INNER JOIN GameUsers ON Games.game_code=GameUsers.game_code AND GameUsers.user_id={message.from_user.id}
    """
    )
    if data_time:
        data_time,title, description = data_time[0]
        date_to_start = (datetime.strptime(data_time,"%d.%m.%Y %H:%M")-datetime.now()).total_seconds()
        date_to_start = f"{int(date_to_start // 3600)} г." if date_to_start >= 3600 else f"{int(date_to_start // 60)} хв."
        await message.answer(f"""
АНГЛіЗі розпочнеться через *{date_to_start}*""", parse_mode='Markdown'
        )
    else:
        await message.answer("Ви не зареєстровані на поточну ігру")

async def contacts(callback : types.CallbackQuery):
    await callback.message.edit_media(
        types.InputMediaPhoto(types.InputFile('Menu/Frame 58.png')),
        markups.contacts_markups
    )
    db.sqlite(f"UPDATE Users SET step = step + 1 WHERE user_id = {callback.from_user.id}")
    step, = db.get_user_info(callback.from_user.id, 'step')
    cell = await create_bot.user_workspace.find(str(callback.from_user.id), in_column=9)
    if cell is not None:
        await create_bot.user_workspace.update_cell(cell.row, config.angli3i_steps.get(callback.data), step)

#                    АРЧІБАЛЬД               #

async def support(callback : types.CallbackQuery):
    await Forms.UserStartForms.Support.set()
    if not db.sqlite(f"SELECT * FROM GameUsers WHERE user_id={callback.from_user.id}"):
        await callback.message.answer_sticker(types.InputFile("stickers/AnimatedSticker.tgs"))
        await callback.message.answer_voice(
            types.InputFile("audios/5307625519590678826.mp4.mp3"),
            reply_markup=markups.back_bttn
        )
    else:
        await callback.message.answer_sticker(types.InputFile("stickers/AnimatedSticker.tgs"))
        await callback.message.answer_voice(
            types.InputFile("audios/5307625519590678841.mp4.mp3"),
            reply_markup=markups.back_bttn
        )

async def exit_support(callback : types.CallbackQuery, state : FSMContext):
    db.sqlite(f"UPDATE Users SET Archibalt = 1 WHERE user_id = {callback.from_user.id}")
    user = db.get_user_info(callback.from_user.id)
    await state.finish()
    #await callback.message.answer_sticker(types.InputFile("stickers/AnimatedSticker.tgs"))
    #await callback.message.answer_voice(open("audios/5307625519590679508.mp4.mp3", 'rb'))

    await callback.message.answer_photo(
        open('Menu/Frame 39.png', 'rb'),
        reply_markup=markups.Menu if user[5] == 0 else markups.MainMenu
    )

#                MY CHAT MEMBER                #

async def user_block(UpdateData : types.ChatMemberUpdated, state : FSMContext):
    await state.finish()
    cell = await create_bot.user_workspace.find(str(UpdateData.from_user.id), in_column=9)
    if cell is not None:
        await create_bot.user_workspace.update_cell(cell.row, 6,'Ні')

async def user_start(UpdateData : types.ChatMemberUpdated):
    cell = await create_bot.user_workspace.find(str(UpdateData.from_user.id), in_column=9)

    if cell is not None:
        await create_bot.user_workspace.update_cell(cell.row, 6,'Так')

async def user_joined_chat(message: types.Message):
        print(f"NEW CHAT MEMBER {message.from_user.username}")
        try:
            num_played_games, = db.get_user_info(message.from_user.id, 'num_played_games')
            if num_played_games < 1 or not db.sqlite(f"SELECT * FROM GameUsers WHERE user_id = {message.from_user.id}"):
                await bot.kick_chat_member(-1001744110388, message.from_user.id)
                await bot.unban_chat_member(-1001744110388, message.from_user.id)
        except Exception:
            print("Some error", message.from_user.id, message.from_user.username)
            await bot.kick_chat_member(-1001744110388, message.from_user.id)
            await bot.unban_chat_member(-1001744110388, message.from_user.id)

async def checkout_process(pre_checkout_query : types.PreCheckoutQuery):
    await bot.answer_pre_checkout_query(pre_checkout_query.id, ok=True)

async def successful_payment(message : types.Message):
    game_code = message.successful_payment.invoice_payload.split('_')[-1]
    game_code = game_code if game_code != "None" else None
    level, num_played_games = db.get_user_info(message.from_user.id, 'level, num_played_games')
    db.reg_player(message.from_user.id, message.from_user.username, message.from_user.full_name, game_code, level)
    if num_played_games > 0:
        await message.answer_animation(types.InputFile("images/see you soon.mp4"))
    await message.answer_video_note(
        types.InputFile('WelcomePageItems/doc_2023-01-03_18-45-43.mp4')
    )
    await message.answer_photo(
        open('Menu/Frame 39.png', 'rb'),
        reply_markup=markups.MainMenu
    )
    cell = await create_bot.user_workspace.find(str(message.from_user.id), in_column=10)
    if cell:
        total_amount, = db.get_user_info(message.from_user.id, 'total_amount')
        await create_bot.user_workspace.update_acell(f'G{cell.row}', total_amount+message.successful_payment.total_amount)
        db.sqlite(f"UPDATE Users SET total_amount = total_amount+{message.successful_payment.total_amount} WHERE user_id = {message.from_user.id}")

#                REGISTER HANDLERS             #

def register_user_handlers(dp : Dispatcher):

    dp.register_message_handler(start, commands=['start'])

    dp.register_message_handler(user_joined_chat,lambda message:message.chat.id == -1001744110388 and message.from_user.id != bot.id,content_types=[types.ContentType.NEW_CHAT_MEMBERS, types.ContentType.LEFT_CHAT_MEMBER])

    dp.register_callback_query_handler(start_menu, text='back_to_menu')
    dp.register_callback_query_handler(reg_welcome_page, text="register")
    dp.register_callback_query_handler(reg_game, text_contains="reg_game_")
    dp.register_callback_query_handler(angli3i_info, text=['open_details_angli3i', 'menu_back'])
    dp.register_callback_query_handler(open_menu_angli3i,text=['menu_whoplays','menu_wheregoing','menu_howgoing', 'menu_price'])
    dp.register_callback_query_handler(contacts, text='contacts')

    #dp.register_message_handler(time_to_game,text="СКІЛЬКИ ЧАСУ ДО АНГЛіЗі", state=[None, Forms.UserStartForms.user, Forms.UserStartForms.user1])
    dp.register_callback_query_handler(support, text="connect_help")
    dp.register_callback_query_handler(exit_support, text='back_to_menu', state=Forms.UserStartForms.Support)

    dp.register_my_chat_member_handler(user_block,lambda i:i.new_chat_member.status == "kicked", state=[None, '*'])
    dp.register_my_chat_member_handler(user_start)
    dp.register_pre_checkout_query_handler(checkout_process)
    dp.register_message_handler(successful_payment,content_types=[types.ContentType.SUCCESSFUL_PAYMENT])


import csv
import os
import numpy as np
import matplotlib.pyplot as plt

# Function for adding a new student to the records
def add_student(students):
    # Prompt user for input data
    student_number = input("Enter student number: ")
    surname = input("Enter surname: ")
    given_name = input("Enter given name: ")
    unit_mark = float(input("Enter unit mark: "))

    # Create dictionary of student data and append it to the array of students
    student = {
        'student_number': student_number,
        'surname': surname,
        'given_name': given_name,
        'unit_mark': unit_mark
    }
    students = np.append(students, student)

    # Print success message and return updated array of students
    print("Student details added successfully.")
    return students

# Function for searching the records for students based on a keyword
def search_students(students):
    # Prompt user for search keyword
    search_key = input("Enter student number or name to search: ").lower()
    matching_students = []

    # Check each student in the array for a match with the search key
    for student in students:
        if (search_key in student['student_number'].lower()) or (search_key in student['surname'].lower()) or (search_key in student['given_name'].lower()):
            matching_students.append(student)

    # Print list of matching students if any found, else print error message
    if len(matching_students) > 0:
        print("Matching students found:")
        for student in matching_students:
            print_student_details(student)
    else:
        print("No matching students found.")

# Helper function for printing details of a single student record
def print_student_details(student):
    print("Student Number:", student['student_number'])
    print("Surname:", student['surname'])
    print("Given Name:", student['given_name'])
    print("Unit Mark:", student['unit_mark'])
    print()

# Function for displaying all students in a particular grade range
def display_students_by_grade(students):
    # Prompt user for target grade
    grade = input("Enter unit grade (HD, D, C, P, or N): ")
    grade_students = []

    # Check each student in the array for a match with the target grade
    for student in students:
        student_grade = convert_mark_to_grade(student['unit_mark'])
        if student_grade == grade:
            grade_students.append(student)

    # Print list of matching students if any found, else print error message
    if len(grade_students) > 0:
        print("Students with grade", grade + ":")
        for student in grade_students:
            print_student_details(student)
    else:
        print("No students found with grade", grade + ".")

# Helper function for converting a numeric mark to a letter grade
def convert_mark_to_grade(mark):
    if mark >= 85:
        return "HD"
    elif mark >= 75:
        return "D"
    elif mark >= 65:
        return "C"
    elif mark >= 50:
        return "P"
    else:
        return "N"

# Function for deleting a student record from the array
def delete_student(students):
    # Prompt user for target student number
    student_number = input("Enter student number to delete: ")
    deleted = False

    # Check each student in the array for a match with the target student number
    for index, student in enumerate(students):
        if student['student_number'] == student_number:
            students = np.delete(students, index)
            deleted = True
            break

    # Print success message or error message if no matching student found
    if deleted:
        print("Student", student_number, "deleted successfully.")
    else:
        print("Student", student_number, "not found.")

    # Return updated array of students
    return students

# Function for loading student records from a CSV file and adding them to the existing array
def load_students_from_file(students):
    # Prompt user for file path and check if file exists
    file_path = input("Enter path to CSV file: ")
    if not os.path.exists(file_path):
        print("File not found:", file_path)
        return students

    # Open file and read each row as a new student, then add to the array
    with open(file_path, 'r', newline='') as csvfile:
        reader = csv.reader(csvfile)
        next(reader)  # skip header row
        for row in reader:
            try:
                student_number = row[0]
                surname = row[1]
                given_name = row[2]
                unit_mark = float(row[3])
                student = {
                    'student_number': student_number,
                    'surname': surname,
                    'given_name': given_name,
                    'unit_mark': unit_mark
                }
                students = np.append(students, student)
            except ValueError:
                print("Error loading row:", row)

        # Print success message and return updated array of students
        print(len(students), "students loaded from file", file_path)
    return students

# Function for saving all current student records to a CSV file
def save_students_to_file(students):
    # Prompt user for file path and check if file exists, giving option to overwrite if it does
    file_path = input("Enter path to save CSV file: ")
    if os.path.exists(file_path):
        choice = input(f"{file_path} already exists. Do you want to overwrite it? (y/n)")
        if choice != 'y':
            print("Operation cancelled. No changes were made.")
            return

    # Open file and write header row, then write each student record as a new row
    with open(file_path, 'w', newline='') as csvfile:
        fieldnames = ['student_number', 'surname', 'given_name', 'unit_mark']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for student in students:
            writer.writerow(student)
        print(len(students), "students saved to file", file_path)

# Function for displaying a pie chart showing the distribution of grades among all students
def display_grade_distribution(students):
    grades, counts = np.unique([convert_mark_to_grade(student['unit_mark']) for student in students], return_counts=True)
    plt.pie(counts, labels=grades)
    plt.title('Grade Distribution')
    plt.show()

# Function for displaying a bar chart showing the distribution of numeric marks among all students
def display_marks_distribution(students):
    marks = [student['unit_mark'] for student in students]
    bins = [0, 30, 40, 50, 60, 70, 80, 90, 100]
    categories = ['0-29', '30-39', '40-49', '50-59', '60-69', '70-79', '80-89', '90-100']
    hist, _ = np.histogram(marks, bins=bins)
    percent = hist/len(marks) * 100
    plt.bar(categories, percent)
    plt.title('Marks Distribution')
    plt.xlabel('Mark Ranges')
    plt.ylabel('Percentage')
    plt.show()

# Function for printing the main menu of options and prompting user for choice
def print_menu():
    print("Menu:")
    print("1. Add student details")
    print("2. Search for students")
    print("3. Display students by grade")
    print("4. Delete student record")
    print("5. Load student records from a CSV file")
    print("6. Save the student records in the current system to a CSV file")
    print("7. Quit")
    print("8. Display grade distribution")
    print("9. Display marks distribution")

# Main function that controls program flow
def main():
    students = np.array([])

    # Load initial student records from file
    students = load_students_from_file(students)

    # Display main menu and prompt user for choice until they choose to quit
    while True:
        print_menu()
        choice = input("Enter your choice (1-9): ")
        if choice == '1':  # Call add_student() function if user chooses option 1
            students = add_student(students)
        elif choice == '2':  # Call search_students() function if user chooses option 2
            search_students(students)
        elif choice == '3':  # Call display_students_by_grade() function if user chooses option 3
            display_students_by_grade(students)
        elif choice == '4':  # Call delete_student() function if user chooses option 4
            students = delete_student(students)
        elif choice == '5':  # Call load_students_from_file() function if user chooses option 5
            students = load_students_from_file(students)
        elif choice == '6':  # Call save_students_to_file() function if user chooses option 6
            save_students_to_file(students)
        elif choice == '7':  # End program if user chooses option 7
            print("Quitting the program...")
            break
        elif choice == '8':  # Call display_grade_distribution() function if user chooses option 8
            display_grade_distribution(students)
        elif choice == '9':  # Call display_marks_distribution() function if user chooses option 9
            display_marks_distribution(students)
        else:
            # Print error message and re-prompt user for menu choice if they entered an invalid option
            print("Invalid choice. Please try again.")

if __name__ == '__main__':
    main()


from datetime import datetime
import time
import webbrowser

import httpx

from client_refresh_flow import RefreshFlow

def login(client):
    info = client.post("/authorize").json()
    print(f"""You have {info['expires_in']} seconds to enter the code

    {info['user_code']}

after authorizing at the URL:

    {info['authorization_uri']}
""")
    webbrowser.open(info["authorization_uri"])  # Try to open a web browser.
    deadline = datetime.now().timestamp() + info["expires_in"]

    # Poll the /token endpoing waiting for this pending session to be verified.
    print("Waiting...", end="", flush=True)
    while datetime.now().timestamp() < deadline:
        response = client.post("/token", data={"device_code": info["device_code"]})
        print(".", end="", flush=True)
        if response.status_code == 200:
            print("\nLogged in!")
            break
        time.sleep(info["interval"])
    tokens = response.json()
    client.auth = RefreshFlow(tokens, f"{client.base_url}/refresh")
    return client

if __name__ == "__main__":
    import time

    client = httpx.Client(base_url="http://localhost:8000")
    login(client)
    print(client.get("/data"))
    print(client.get("/data").json())


import io
from mnist.model import LogisticRegression
import boto3
import numpy as np
import os
import torch
from mnist.configuration import Configuration
from PIL import Image
import torchvision.transforms as transforms

def download_model_artifacts(config: Configuration):
    """
    Downloads model artifact from S3
    """
    model_name, model_path = config.model_artifact.split("/")[-1], config.model_artifact

    boto3_client = boto3.client(
        "s3",
        region_name=config.s3["region_name"],
        aws_access_key_id=config.s3["access_key"],
        aws_secret_access_key=config.s3["secret_key"],
    )
    boto3_client.download_file(config.s3["bucket"], model_name, model_path)

def load_model(config: Configuration):

    if not os.path.exists(config.model_artifact):
        download_model_artifacts(config)

    model = LogisticRegression()
    model.load_state_dict(torch.load(config.model_artifact))
    model.eval()
    return model

def transform_image(image_bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("L")
    image = np.asarray(transforms.Resize(28)(image))
    image = image.reshape(784) / 255
    return torch.Tensor(image)


#Get Details from Spotify API using Spotipy.
# Details we are trying to get: 1. Playlist of the user. 2. Songs in Playlist 3. Analysis of Songs in the playlist.#

import os
import spotipy
import webbrowser
import json,csv
import pandas as pd
from json.decoder import JSONDecodeError
import spotipy.util as util
from spotipy.oauth2 import SpotifyOAuth
from dotenv import load_dotenv
#Importing Secure Keys using ENV (dotenv) Module
load_dotenv()

#User Credentials
client_id=os.getenv('SPOTIFY_CLIENT_ID')
client_secret=os.getenv('SPOTIFY_CLIENT_SECRET')
redirect_uri=os.getenv('REDIRECT_URI')
#Spotify Token Authentication
scope = os.getenv('SCOPES')

final_extract=".\\final_data\\extract_data\\"
final_stream='.\\final_data\\StreamingHistory\\'
spotify_data='.\\spotify_data\\vijayasai\\'

sp = spotipy.Spotify(auth_manager=SpotifyOAuth(client_id=client_id,client_secret=client_secret,
                                                redirect_uri=redirect_uri, scope=scope))

# # Prints the Current User's Details
# user=sp.current_user()
# print(json.dumps(user,indent=2))

#Gets Current Song played by the user

current_song=sp.current_user_playing_track()
print(current_song)
if(current_song==None):
    print("No Song is Playing Right Now! Please check your Spotify Again!!")
else:
    with open('current_song.json','w') as current_song_output:
        json.dump(current_song,current_song_output)
    print("Name of Current Song: "+current_song['item']['name']+
        "\nURI of the Song: "+current_song['item']['uri'])
    uri_current_song = current_song['item']['uri']
    audio_analysis= sp.audio_features(tracks=uri_current_song)
    # with open(os.path.join(final_extract,'current_song_analysis.json'),'w') as current_song_output:
    #     json.dump(audio_analysis,current_song_output)
    # print("Audio Analysis for current song extracted successfully")

current_song_cover=current_song["item"]["album"]["images"][0]["url"]
print(current_song_cover)

# Getting top tracks of the user
# top_tracks=sp.current_user_top_tracks(limit=20, offset=0, time_range='long_term')
# with open(os.path.join(final_extract,'top_tracks_long.json'),'w') as current_song_output:
#     json.dump(top_tracks, current_song_output)
# top_tracks_json=[]
# for i in range(len(top_tracks["items"])):
#     top_tracks_json.append(dict(trackName=top_tracks["items"][i]["name"],
#                                     trackURI=top_tracks["items"][i]["uri"]))
#     with open(os.path.join(final_extract,'top_tracks.json'),'w') as outfile:
#         json.dump(top_tracks_json, outfile)

# results = sp.current_user_saved_tracks()
# for idx, item in enumerate(results['items']):
#     track = item['track']
#     print(idx, track['artists'][0]['name'], " – ", track['name']," - ", track['id'])
# list_playlists=sp.current_user_playlists(limit=50, offset=0)
# print(json.dumps(list_playlists,indent=4))

# items_in_playlists=sp.playlist_items('2XIivRPHOyhKkm8FkphG9V', fields="items(track(name,id))", limit=100, offset=0, market=None, additional_types=('track', 'episode'))
# print(type(items_in_playlists))
# print(json.dumps(items_in_playlists,indent=4))
# print(items_in_playlists["items"][i]["track"]["id"])
# print(sp.audio_features(tracks=[items_in_playlists["items"][i]["track"]["id"]]),fields="danceability,energy,loudness,speechiness,acousticness")


from collections import defaultdict

memo = defaultdict(int)

red_tile = 2
green_tile = 3
blue_tile = 4

def calculate(place):
    if red_tile > place:
        return 1
    if memo[place]:
        return memo[place]

    place_red = calculate(place - red_tile)
    place_green, place_blue = 0, 0

    if green_tile <= place:
        place_green = calculate(place - green_tile)
    if blue_tile <= place:
        place_blue = calculate(place - blue_tile)

    skip_square = calculate(place - 1)

    r = place_red + place_green + place_blue + skip_square
    memo[place] = r
    return r

print(calculate(50))


import logging
from collections import OrderedDict
import numpy as np
import torch
import cv2
import math
import torch.nn as nn
from torch.nn.parallel import DataParallel, DistributedDataParallel
import models.networks as networks
import models.lr_scheduler as lr_scheduler
from .base_model import BaseModel
from models.modules.loss import ReconstructionLoss
from models.modules.ffmpeg_DE import DecodeEncode, DownSample, UpSample #codecImage, rgbdown
import sys
import torch.nn.functional as F
from models.modules.Quant import BypassRound, Quant
from models.modules.format_trans import rgb2yuv420_down, yuv4202rgb

logger = logging.getLogger('base')

class IRNModel(BaseModel):
    def __init__(self, opt):
        super(IRNModel, self).__init__(opt)

        if opt['dist']:
            self.rank = torch.distributed.get_rank()
        else:
            self.rank = -1  # non dist training
        train_opt = opt['train']
        test_opt = opt['test']
        self.train_opt = train_opt
        self.test_opt = test_opt
        self.scalea = opt['scalea']
        self.scaleb = opt['scaleb']

        self.netG = networks.define_G(opt).to(self.device)
        self.netA = networks.define_Analysis(opt).to(self.device)
        self.netS = networks.define_S(opt).to(self.device)

        if opt['dist']:
            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
            self.netA = DistributedDataParallel(self.netA, device_ids=[torch.cuda.current_device()])
            self.netS = DistributedDataParallel(self.netS, device_ids=[torch.cuda.current_device()])

        else:
            self.netA = DataParallel(self.netA)
            self.netG = DataParallel(self.netG)
            self.netS = DataParallel(self.netS)

        # print network
        self.print_network()
        self.load()

        self.DecodeEncode = DecodeEncode()
        self.bypass_round = BypassRound

        if self.is_train:
            self.netG.train()
            self.netA.train()
            self.netS.train()

            # loss
            self.Reconstruction_forw = ReconstructionLoss(losstype=self.train_opt['pixel_criterion_forw'])
            self.Reconstruction_back = ReconstructionLoss(losstype=self.train_opt['pixel_criterion_back'])

            # optimizers
            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
            optim_params = []
            for k, v in self.netG.named_parameters():
                if v.requires_grad:
                    optim_params.append(v)
                else:
                    if self.rank <= 0:
                        logger.warning('Params [{:s}] will not optimize.'.format(k))
            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], weight_decay=wd_G, betas=(train_opt['beta1'], train_opt['beta2']))
            self.optimizers.append(self.optimizer_G)

            optim_params = []
            optim_rates = []
            for k, v in self.netA.named_parameters():
                if v.requires_grad:
                    if ('rate' not in k):
                        optim_params.append(v)
                    else:
                        optim_rates.append(v)
                else:
                    if self.rank <= 0:
                        logger.warning('Params [{:s}] will not optimize.'.format(k))
            self.optimizer_A = torch.optim.Adam(optim_params, lr=train_opt['lr_A'], weight_decay=wd_G, betas=(train_opt['beta1'], train_opt['beta2']))
            self.optimizers.append(self.optimizer_A)

            #self.optimizer_C = torch.optim.Adam(optim_rates, lr=train_opt['lr_A'], weight_decay=wd_G, betas=(train_opt['beta1'], train_opt['beta2']))
            #self.optimizers.append(self.optimizer_C)

            optim_params = []
            for k, v in self.netS.named_parameters():
                if v.requires_grad:
                    optim_params.append(v)
                else:
                    if self.rank <= 0:
                        logger.warning('Params [{:s}] will not optimize.'.format(k))
            self.optimizer_S = torch.optim.Adam(optim_params, lr=train_opt['lr_S'], weight_decay=wd_G, betas=(train_opt['beta1'], train_opt['beta2']))
            self.optimizers.append(self.optimizer_S)

            # schedulers
            if train_opt['lr_scheme'] == 'MultiStepLR':
                for optimizer in self.optimizers:
                    self.schedulers.append(
                        lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
                                                         restarts=train_opt['restarts'],
                                                         weights=train_opt['restart_weights'],
                                                         gamma=train_opt['lr_gamma'],
                                                         clear_state=train_opt['clear_state']))
            elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
                for optimizer in self.optimizers:
                    self.schedulers.append(
                        lr_scheduler.CosineAnnealingLR_Restart(
                            optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
                            restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
            else:
                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')

            self.log_dict = OrderedDict()

    def gaussian_batch(self, dims):
        return torch.randn(tuple(dims)).to(self.device)

    def loss_forward(self, out, y):
        l_forw_fit = self.train_opt['lambda_fit_forw'] * self.Reconstruction_forw(out, y)

        return l_forw_fit

    def loss_backward(self, out, y):
        l_back_rec = self.train_opt['lambda_rec_back'] * self.Reconstruction_back(out, y)

        return l_back_rec

    def feed_data(self, data):
        #self.ref_L_rgb = data['LQ'].to(self.device)  # LQ
        self.real_H_rgb = data['GT'].to(self.device)  # GT

    def optimize_parameters(self, step):

        ### compress
        qp = "6"
        print('real_H_rgb size: ', self.real_H_rgb.shape)
        self.real_H, self.ref_L = rgb2yuv420_down(self.real_H_rgb, self.scalea, self.scaleb) # (B * 6 * H/2 * W/2 , B * 6 * H/4 * W/4)
        print('real_H size: ', self.real_H.shape)
        print('ref_L size: ', self.ref_L.shape)
        _, _, h, w = self.real_H.size()
        out_mv = None

        # update netA virtual codec
        lr_lanczos = self.ref_L # B * 6 * H/4 * W/4
        lr_lanczos_codec, out_mv = self.DecodeEncode(lr_lanczos, qp, False)
        self.optimizer_A.zero_grad()
        lr_lanczos_codec_est, _ = self.netA(x=lr_lanczos, mv=out_mv, I=True)
        l_codec_est = self.loss_forward(lr_lanczos_codec_est, lr_lanczos_codec) # loss between netA & x265 decoded lr
        loss_A = l_codec_est
        loss_A.backward()
        torch.nn.utils.clip_grad_norm_(self.netA.parameters(), 1.0)
        self.optimizer_A.step()

        # update encoder & decoder netG (& rate estimation network in netA), netS
        self.optimizer_G.zero_grad()
        # self.optimizer_C.zero_grad()
        y, y_syntax, content = self.netG(x=self.real_H, up=False)  # forward downscaling, B * 6 * H/4 * W/4
        print('content size: ', content.shape) # B * 6 * H/4 * W/4
        print('y size: ', y.shape) # B * 192 * H/32 * W/32
        print('y_syntax size: ', y_syntax.shape) # B * 16 * H/32 * W/32
        content_codec_est, bpp_netA = self.netA(x=content, mv=out_mv, I=True)
        # l_size_fit = torch.mean(bpp_netA)
        sr_netG = self.netG(x=content_codec_est, up=True) # B * M * H/2 * W/2

        # Net for syntax
        bpp_netS, sr_tilde = self.netS(y = y, x_tilde = sr_netG)  # B * 6 * H/2 * H/2

        l_tilde_real = self.loss_forward(sr_tilde, self.real_H)
        l_content_ref = self.loss_forward(content, self.ref_L)

        loss = l_tilde_real + 0.5 * l_content_ref  + bpp_netS #+ 8e-4 * l_size_fit
        loss = loss.mean()
        l_tilde_real = l_tilde_real.mean()
        l_content_ref = l_content_ref.mean()
        bpp_netS = bpp_netS.mean()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.netG.parameters(), 1.0)
        torch.nn.utils.clip_grad_norm_(self.netS.parameters(), 1.0)
        self.optimizer_G.step()
        self.optimizer_S.step()
        #self.optimizer_C.step()

        # set log
        self.log_dict['l_codec_est'] = l_codec_est.item()
        self.log_dict['l_content_ref'] = l_content_ref.item()
        self.log_dict['l_tilde_real'] = l_tilde_real.item()
        self.log_dict['bpp_netS'] = bpp_netS.item()
        self.log_dict['distortion'] = (l_tilde_real + 0.5 * l_content_ref).item()
        #self.log_dict['l_size_fit'] = l_size_fit.item()

    def test(self):
        qp = "22"
        self.netG.eval()
        self.netS.eval()
        with torch.no_grad():
            # YUV train
            real_H, ref_L = rgb2yuv420_down(self.real_H_rgb, self.scalea, self.scaleb) # (B * 6 * H/2 * W/2 , B * 6 * H/4 * W/4)
            self.ref_L = yuv4202rgb(ref_L) # B * 3 * H/2 * W/2
            # _, _, h, w = real_H.size()
            # ref_L, _ = self.DecodeEncode(ref_L, qp, True) # B * 6 * H/4 * W/4
            # fake_H_bic = self.yuvUpsample(ref_L, h, w) # b * 6 * height//2 * width//2
            # self.fake_H_bic = yuv4202rgb(fake_H_bic) # b * 3 * height * width
            # self.psnr_fix = 10 * torch.log10(1**2 / torch.mean((real_H - fake_H_bic)**2))

            y, y_syntax, content = self.netG(x=real_H, up=False) # forward downscaling, B * 6 * H/4 * W/4
            self.forw_L = yuv4202rgb(content) # B * 3 * H/2 * W/2
            content_codec, _ = self.DecodeEncode(content, qp, True) # B * 6 * H/4 * W/4
            sr_netG = self.netG(x=content_codec, up=True) # B * M * H/2 * W/2
            bpp_netS, sr_tilde = self.netS(y = y, x_tilde = sr_netG, mode='test')  # B * 6 * H/2 * H/2
            self.fake_H = yuv4202rgb(sr_tilde) # B * 3 * H * W
            self.psnr = 10 * torch.log10(1**2 / torch.mean((real_H - sr_tilde)**2))
            #self.psnr = 10 * torch.log10(255**2 / torch.mean((real_H - sr_tilde)**2))
            self.bpp_netS = bpp_netS
            self.psnr = self.psnr.mean()
            self.bpp_netS = self.bpp_netS.mean()

        self.netG.train()
        self.netS.train()

    def downscale(self, HR_img):
        self.netG.eval()
        with torch.no_grad():
            y, y_syntax, LR_img = self.netG(x=HR_img, up=False)
        return LR_img

    def upscale(self, LR_img, scale, gaussian_scale=1):
        with torch.no_grad():
            HR_img = self.netG(x=LR_img, up=True, flow=None, I=True)
        self.netG.eval()
        return HR_img

    def get_current_log(self):
        return self.log_dict

    def get_current_visuals(self):
        out_dict = OrderedDict()
        out_dict['LR_ref'] = self.ref_L.detach()[0].float().cpu()
        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
        out_dict['LR'] = self.forw_L.detach()[0].float().cpu()
        out_dict['GT'] = self.real_H.detach()[0].float().cpu()
        #out_dict['SR_bic'] = self.fake_H_bic.detach()[0].float().cpu()
        #out_dict['PSNR_fix'] = self.psnr_fix.detach().item()
        out_dict['PSNR'] = self.psnr.detach().item()
        out_dict['bpp_netS'] = self.bpp_netS.detach().item()
        # out_dict['SSIM_fix'] = self.ssim_fix.detach().item()
        # out_dict['SSIM'] = self.ssim.detach().item()
        return out_dict

    def print_network(self):
        s, n = self.get_network_description(self.netG)
        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
                                             self.netG.module.__class__.__name__)
        else:
            net_struc_str = '{}'.format(self.netG.__class__.__name__)
        if self.rank <= 0:
            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
            logger.info(s)

    def load(self):
        load_path_G = self.opt['path']['pretrain_model_G']
        if load_path_G is not None:
            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])

        load_path_A = self.opt['path']['pretrain_model_A']
        if load_path_A is not None:
            logger.info('Loading model for A [{:s}] ...'.format(load_path_A))
            self.load_network(load_path_A, self.netA, self.opt['path']['strict_load'])

        load_path_S = self.opt['path']['pretrain_model_S']
        if load_path_S is not None:
            logger.info('Loading model for S [{:s}] ...'.format(load_path_S))
            self.load_network(load_path_S, self.netS, self.opt['path']['strict_load'])

    def save(self, iter_label):
        self.save_network(self.netG, 'G', iter_label)
        self.save_network(self.netA, 'A', iter_label)
        self.save_network(self.netS, 'S', iter_label)


"""
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()


from django.contrib import admin

# Register your models here.
from .models import Transaction, TransactionDetails

class DetailInline(admin.TabularInline):
    model = TransactionDetails

class TransactionAdmin(admin.ModelAdmin):
    list_display = ('shopper', 'total_price', 'payment_id', 'is_valid')
    inlines = [DetailInline]

admin.site.register(Transaction, TransactionAdmin)


import matplotlib.pyplot as plt

f = open("Utility_vs_Runtime.txt","r")
y = []
x = []
param = []
line = []
lines = f.readlines()
for i in lines:
	line = i.split(",")
	param.append(float(line[0]))
	x.append(round(float(line[1]),3))
	y.append(round(float(line[2][:-1]),1))
f.close()

fig, (ax1,ax2) = plt.subplots(nrows =1, ncols =2)
ax1.set_title("Sub-sampling frames: Utility runtime tradeoff")
ax1.set_ylabel("Run-time")
ax1.set_xlabel("Utility")
ax1.scatter(x, y, label='Data')
ax1.plot(x, y, linestyle = 'dashed')
for i in range(len(x)):
	ax1.annotate(" (" + str(x[i]) + "," + str(y[i]) + ")" ,(x[i], y[i]), horizontalalignment='right')
ax1.grid()

ax2.set_title("Sub-sampling frames: Runtime vs Sub-sampling parameter")
ax2.set_ylabel("Run-time")
ax2.set_xlabel("Sub-sampling parameter")
ax2.scatter(param,y, label='Data')
ax2.plot(param,y, linestyle = 'dashed')
for i in range(len(y)):
	ax2.annotate(" (" + str(param[i]) + "," + str(y[i]) + ")" ,(param[i], y[i]), horizontalalignment='right')
ax2.grid()
plt.tight_layout()
plt.show()


import inspect

from marvin.bots.utilities import UtilityBot
from marvin.programs import Program

class ApproximatelyEquivalent(Program):
    """
    Determines if two statements are approximately (semantically) equivalent.
    """

    async def run(self, statement_1: str, statement_2: str):
        bot = UtilityBot(instructions=inspect.cleandoc("""
                The user will give you two statements. Your only job is to
                determine if the two statements are approximately equivalent.
                You will respond ONLY with either the word `true` or `false`. Do
                not say anything else, for any reason.
                """))

        message = inspect.cleandoc("""
            # Statement 1
            {statement_1}

            # Statement 2
            {statement_2}
            """).format(statement_1=statement_1, statement_2=statement_2)

        response = await bot.say(message)
        if response.content.lower() == "true":
            return True
        else:
            return False