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

""" This script collects all the data in orgs and sources folders and merge them in a single json file. """
import json, pathlib, os, sys

### ENVIRONMENTAL VARIABLES
# environmental variables can be set in order to override default values
# NOTE: you can use relative or absolute paths, with or without a separator at the end of folder names

# the folder that contains sources json files
# default: './sources'
env_sources = 'OPENDATA_SOURCES_DIR'

# the folder containing the organization details
# default: './orgs'
env_organizations = 'OPENDATA_ORGANIZATIONS_DIR'

# the filename that will store all results (include extension)
# default: './dist/index.json'
env_dist_filename = 'OPENDATA_DIST_FILENAME'

# the filename that will store nested results (include extension)
# default: './dist/nested/index.json'
env_nested_filename = 'OPENDATA_NESTED_FILENAME'

# shall the script override the data?
# default: True
env_allow_override = 'OPENDATA_CAN_OVERRIDE'

# It may be desiderable to remove the owner_org key from the source since it is implicit.
# This saves a few bytes in the final json file. If you want to keep the owner_org key
# feel free to set the variable to True
# default: False
env_keep_owner = 'OPENDATA_KEEP_OWNER'

# in case you want just to output to the console (i.e. if you want to pipe the results into a parser)
# default: False
env_to_stdout = 'OPENDATA_USE_STDOUT'

### DEFAULT SETTINGS
falsy_strings = ('no', 'false', 'never', 'n', 'f', 'falso', 'mai') # add other strings if necessary (?)
empty = ('', None)

sources_dir = os.environ[env_sources] if (env_sources in os.environ) and (os.environ[env_sources] not in empty) else pathlib.Path('.', 'sources')
orgs_dir = os.environ[env_organizations] if (env_organizations in os.environ) and (os.environ[env_organizations] not in empty) else pathlib.Path('.',  'orgs')
dist_filename = os.environ[env_dist_filename] if (env_dist_filename in os.environ) and (os.environ[env_dist_filename] not in empty) else  pathlib.Path('.',  'dist/index.json')
nested_filename = os.environ[env_nested_filename] if (env_nested_filename in os.environ) and (os.environ[env_nested_filename] not in empty) else  pathlib.Path('.',  'dist/nested/index.json')
override = (os.environ[env_allow_override].lower() not in falsy_strings) if (env_allow_override in os.environ) and (os.environ[env_allow_override] not in empty) else True
keep_owner = (os.environ[env_keep_owner].lower() not in falsy_strings) if (env_keep_owner in os.environ) and (os.environ[env_keep_owner] not in empty) else False
to_stdout = (os.environ[env_to_stdout].lower() not in falsy_strings) if (env_to_stdout in os.environ) and (os.environ[env_to_stdout] not in empty) else False

# A dictionary to guide in the classification of the organizations.
# There are two main branches, "nazionale" (national) and "locale" (local).
# Every branch has a inner dictionary. The inner dictionary keys are the first word in org.title whereas
# the dictionary values are the keys to be used to identify the type of organization in json output.
# You can customize the values returned; the key "*" is used as a catch-all alternative if the first word
# in org.title is not present in the dictionary's branch.
classification = {
    'nazionale': {
        'ministero': 'ministero',
        '*': 'altro'
    },
    'locale': {
        'citta': 'citta metropolitana',
        'comune': 'comune',
        'provincia': 'provincia',
        'regione': 'regione',
        'universita': 'universita',
        '*': 'altro'
    }
}

### UTILITIES
def classify(organization):
    """
    the function checks the first word in the title of the organization
    and returns a list of keys to be used to classify it.
    """
    first_word = organization['name'].split('-')[0]
    category = 'locale' if 'region' in organization.keys() else 'nazionale'
    result = [category]

    if category == 'locale':
        result.append(organization['region'])

    if first_word in classification[category].keys():
        result.append(classification[category][first_word])
    else:
        result.append(classification[category]['*']) # first word not recognized.

    return result

def populate_dict(keys_list, dictionary, organization, source):
    """
    recursive function that takes a list of keys to be added to a dict of dicts (the dictionary argument).
    If the list is empty, it returns the organization argument (the leaf) otherwise it returns a dictionary
    created from the nested keys (the branches).

    example:
    --------
    keys_list = ['a', 'b', 'c']
    dictionary = {'other':{'nested'}, 'a':{'foo':'bar'}}
    organization = {"whatever": "value", "you":"want"}

    > populate_dict(keys_list, dictionary, organization)
    > {'other':{'nested'}, 'a':{'foo':'bar', 'b':{'c':{"whatever": "value", "you":"want"}}}}

    """
    if len(keys_list) == 0:
        # time to save the new source
        has_organization = False

        if not keep_owner:
            source.pop('owner_org', None)

        # check if organization is already present
        for org in dictionary:
            if org['name'] == organization['name']:
                # the organization already esists
                organization = org
                # if the organization is already in the dictionary the 'sources' key has been set
                # so it is not necessary to check for its existence
                organization['sources'].append(source)
                has_organization = True
                break

        if not has_organization:
            # no organization found or dictionary is empty
            organization['sources'] = [source]
            dictionary.append(organization)

        return dictionary

    key = keys_list.pop(0)
    if key not in dictionary.keys():
        if len(keys_list) == 0:
            dictionary[key] = populate_dict(keys_list, [], organization, source)
        else:
            dictionary[key] = populate_dict(keys_list, {}, organization, source)
    else:
        dictionary[key] = populate_dict(keys_list, dictionary[key], organization, source)
    return dictionary

### PARSER
def parse():
    """
    the main script
    """

    dist_all = {}
    dist_nested = {}

    for source in pathlib.Path(sources_dir).glob('*.json'):
        with source.open('r') as source_file:
            source_content = json.load(source_file)
            if "config" in source_content:
                source_content["config"] = json.loads(source_content["config"])
            owner = source_content['owner_org']
            try:
                with pathlib.Path(orgs_dir, owner+'.json').open('r') as organization:
                    org_content = json.load(organization)
                    category = classify(org_content)
                    dist_nested = populate_dict(category, dist_nested, org_content, source_content)
                    dist_all[owner] = dist_all.get(owner, dict(org_content, sources=[]))
                    dist_all[owner]["sources"].append({ k:source_content[k] for k in source_content if keep_owner or source_content[k] != 'owner_org' })
            except FileNotFoundError:
                print(f"ERROR: file {pathlib.Path(orgs_dir, owner+'.json')} not found or not readable.", file=sys.stderr)
                exit(2)

    if not dist_nested or not dist_all:
        print(f"WARNING: no sources found. Is {pathlib.Path(sources_dir)} the correct folder?", file=sys.stderr)

    if to_stdout:
        print(json.dumps(dist_all.values(), sort_keys=True, indent=4))

    if override or not os.path.exists(dist_filename):
        with open(dist_filename, 'w') as output_file:
            json.dump(list(dist_all.values()), output_file)
    else:
        print("ERROR: output file exists and I'm not allowed to overwrite it.", file=sys.stderr)

    if override or not os.path.exists(nested_filename):
        with open(nested_filename, 'w') as output_file:
            json.dump(dist_nested, output_file)
    else:
        print("ERROR: output file exists and I'm not allowed to overwrite it.", file=sys.stderr)

### THE SCRIPT
if __name__ == '__main__':
    parse()


"""
Problem:
    1. Two Sum
Difficulty:
    Easy
URL:
    https://leetcode.com/problems/two-sum
Tags:
    Array, Hash Table
Date:
    2022-05-10T14:00:29.877163+08:00
"""
from typing import List

class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        for i, num in enumerate(nums):
            if target - num in nums[i + 1:]:
                return [i, nums.index(target - num, i + 1)]

tests = [
    (
        ([2, 7, 11, 15], 9,
         ),
        [0, 1],
    ),
    (
        ([3, 2, 4], 6,
         ),
        [1, 2],
    ),
    (
        ([3, 3], 6,
         ),
        [0, 1],
    ),
]


import json
import requests  # see http://python-requests.org

def url_for(endpoint):
    return 'http://localhost:5000/{}/'.format(endpoint)

def delete_all_people():
    r = requests.delete(url_for('people'))
    print("'people' deleted, server response:", r.status_code)

def post_people():
    data = [
        {'firstname': 'John', 'lastname': 'Doe'},
        {'firstname': 'Mike', 'lastname': 'Green'},
    ]

    response = requests.post(
        url_for('people'),
        json.dumps(data),
        headers={'Content-Type': 'application/json'}
    )
    print("'people' posted, server response:", response.status_code)

def get_people():
    r = requests.get(url_for('people'))
    print('people downloaded, server response:', r.status_code)

    if r.status_code == 200:
        people = r.json()['_items']
        print('{} people:'.format(len(people)))
        for person in people:
            print('{}, {}'.format(person['firstname'], person['_id']))

def main():
    delete_all_people()
    post_people()
    get_people()

if __name__ == '__main__':
    main()


import numpy as np
import statistics
from scipy import stats

dataset= [5,6,7,5,6,5,7,4,5,5,5,5,7,5,6,6,7,6,6,7,7,7,6,5,6]

#mean value
mean= np.mean(dataset)

#median value
median = np.median(dataset)

#mode value
mode= stats.mode(dataset)

#standard Deviation
Std = statistics.stdev(dataset)

#Variance
Var = statistics.variance(dataset)

print("Mean: ", mean)
print("Median: ", median)
print("Mode: ", mode)
print("Std", Std)
print("Var", Var)


import sqlite3
import sys
import datetime

from collections import defaultdict
from stats_ui_window import Ui_StatWindow
from PyQt5 import QtCore, QtGui, QtWidgets

class MainWindow_EXEC():

    def __init__(self):
        app = QtWidgets.QApplication(sys.argv)
        MainWindow = QtWidgets.QMainWindow()
        self.ui = Ui_StatWindow()
        self.ui.setupUi(MainWindow)

        with open('examiners_names.txt','r') as examiners:
            for line in examiners.readlines():
                self.ui.comboBox.addItem(line.strip())

        self.device_list = []
        self.ui.pushButton.clicked.connect(self.add_device)
        self.ui.pushButton_4.clicked.connect(self.remove_record)
        self.ui.pushButton_2.clicked.connect(self.add_list)
        self.ui.pushButton_3.clicked.connect(QtCore.QCoreApplication.instance().quit)

        MainWindow.show()
        sys.exit(app.exec_())

    def add_device(self):
        device_values = defaultdict()
        device_values['case_number'] = self.ui.lineEdit.text()
        device_values['item_number'] = self.ui.lineEdit_2.text()
        device_values['manufacture'] = self.ui.lineEdit_3.text()
        device_values['model_'] = self.ui.lineEdit_4.text()
        device_values['crime_code'] = self.ui.lineEdit_6.text()
        device_values['requesting'] = self.ui.lineEdit_5.text()
        device_values['examiner'] = str(self.ui.comboBox.currentText())

        if "" in (device_values['case_number'],device_values['item_number'],
        device_values['manufacture'],device_values['model_'],
        device_values['crime_code'],device_values['requesting']):
            self.error_box()
        else:
            all_items = True
            if self.ui.radioButton_11.isChecked():
                device_values['device'] = "Computer"
            elif self.ui.radioButton_10.isChecked():
                device_values['device'] = "Phone"
            elif self.ui.radioButton_12.isChecked():
                device_values['device'] = "Hard Drive"
            elif self.ui.radioButton_13.isChecked():
                device_values['device'] = "Thumbdrive/Media Card"
            elif self.ui.radioButton_14.isChecked():
                device_values['device'] = "Vehilce"
            else:
                all_items = False
                self.error_box(message = "Please Select Device Type")

            if self.ui.radioButton.isChecked():
                device_values['security'] = "Password Protected"
            elif self.ui.radioButton_9.isChecked():
                device_values['security'] = "Unlocked"
            else:
                all_items = False
                self.error_box(message = "Please Select Security")

            if self.ui.checkBox_2.isChecked():
                device_values['secure_start'] = "Enabled"
            else: device_values['secure_start'] = "No"

            if self.ui.checkBox_3.isChecked():
                device_values['logical'] = "Yes"
            else: device_values['logical'] = "No"
            if self.ui.checkBox_4.isChecked():
                device_values['file_system'] = "Yes"
            else: device_values['file_system'] = "No"

            if self.ui.checkBox_5.isChecked():
                device_values['physical'] = "Yes"
            else: device_values['physical'] = "No"

            if self.ui.checkBox_8.isChecked():
                device_values['lt_greykey'] = "Yes"
            else: device_values['lt_greykey'] = "No"

            if self.ui.checkBox_6.isChecked():
                device_values['greykey'] = "Yes"
            else: device_values['greykey'] = "No"

            if self.ui.checkBox_7.isChecked():
                device_values['no_extraction'] = "No Extraction"
            else: device_values['no_extraction'] = "Extracted"

            device_values['date'] = datetime.datetime.now().strftime('%m/%d/%Y')

            if all_items == True:
                self.device_list.append(device_values)

                self.ui.tableWidget.insertRow(0)
                self.ui.tableWidget.setItem(0 , 0, QtWidgets.QTableWidgetItem(device_values['date']))
                self.ui.tableWidget.setItem(0 , 1, QtWidgets.QTableWidgetItem(device_values['device']))
                self.ui.tableWidget.setItem(0 , 2, QtWidgets.QTableWidgetItem(device_values['case_number']))
                self.ui.tableWidget.setItem(0 , 3, QtWidgets.QTableWidgetItem(device_values['item_number']))
                self.ui.tableWidget.setItem(0 , 4, QtWidgets.QTableWidgetItem(device_values['manufacture']))
                self.ui.tableWidget.setItem(0 , 5, QtWidgets.QTableWidgetItem(device_values['model_']))

                self.ui.lineEdit_2.setText("")
                self.ui.lineEdit_3.setText("")
                self.ui.lineEdit_4.setText("")
                self.ui.checkBox_2.setChecked(False)
                self.ui.checkBox_3.setChecked(False)
                self.ui.checkBox_4.setChecked(False)
                self.ui.checkBox_5.setChecked(False)
                self.ui.checkBox_6.setChecked(False)
                self.ui.checkBox_7.setChecked(False)
                self.ui.checkBox_8.setChecked(False)
            else: all_items = True

    def remove_record(self):
        row = self.ui.tableWidget.currentRow()
        self.ui.tableWidget.removeRow(row)

    def add_list(self):
        manufacture = self.ui.lineEdit_3.text()
        if manufacture != "":
            self.error_box(message = "Dont forget to add the phone")
        else:
            self.ui.lineEdit.setText("")
            self.ui.lineEdit_2.setText("")
            self.ui.lineEdit_3.setText("")
            self.ui.lineEdit_4.setText("")
            self.ui.lineEdit_6.setText("")
            self.ui.lineEdit_5.setText("")
            count = self.ui.tableWidget.rowCount()
            if count > 0:
                self.ui.tableWidget.setRowCount(0)
                with open('path.txt','r') as my_path:
                    path = my_path.read()
                con = sqlite3.connect(path)
                cur = con.cursor()
                for item in self.device_list:
                    val = (item['date'],item['case_number'],item['item_number'],item['manufacture'],
                    item['model_'],item['crime_code'],item['requesting'],
                    item['examiner'],item['device'],item['security'],
                    item['secure_start'],item['logical'],item['file_system'],
                    item['physical'],item['lt_greykey'],item['greykey'],item['no_extraction'])
                    sql = "INSERT INTO entries (date,case_number,item_number,manufacture,model_,crime_code,requesting,examiner,device,security,secure_start,logical,file_system,physical,lt_greykey,greykey,no_extraction) VALUES (?,?, ?, ?, ?, ?, ?, ?,?, ?,?,?,?,?,?,?,?)"
                    cur.execute(sql,val)
                    con.commit()
                con.close()

    @staticmethod
    def error_box(message = 'Please fill out all fields!'):
        error_dialog = QtWidgets.QMessageBox()
        error_dialog.setIcon(QtWidgets.QMessageBox.Warning)
        error_dialog.setWindowTitle('Error')
        error_dialog.setText(f'{message}')
        error_dialog.setStandardButtons(QtWidgets.QMessageBox.Close)
        error_dialog.exec()

if __name__ == "__main__":
    MainWindow_EXEC()


# Busca Local Múltiplos Inicios
# Local Search Multiple Starts
import sys
import time

sys.path.insert(1, '../stage_01')
sys.path.insert(1, '../')
from utils import corrent_solution_size, objetive_function, read_instance, viable_solution
from local_search import local_search
from semi_greedy import semi_greedy
import config

# Busca local múltiplos inícios + semi-guloso alpha + busca local primeira melhora
def multiple_starts_local_search(alpha, timeout):
    desks, tests, empty = config.desks, config.tests, config.empty
    desk_count = len(desks)
    test_count = len(tests)
    s_ = viable_solution(desk_count, desk_count, test_count)
    value_ = objetive_function(s_)

    initial_time = time.time()
    current_time = time.time()
    execution_time = current_time - initial_time
    while execution_time < timeout:
        s, _ = semi_greedy(alpha, False)
        s, value = local_search(s, 0, False) # 1 = Primeira  melhora

        if value < value_:
            s_ = s
            value_ = value

        current_time = time.time()
        execution_time = current_time - initial_time

    s_, value_ = corrent_solution_size(s_, empty)
    return s_, value_

if __name__ == '__main__':
    file_name = sys.argv[1]
    timeout = int(sys.argv[2])
    alpha = float(sys.argv[3])
    config.set_timeout(timeout)
    read_instance(file_name)
    s_, value_ = multiple_starts_local_search(alpha, timeout)

    print(s_)
    print(value_)


from django.contrib import admin
from django.urls import path
from . import views

urlpatterns = [
    path('api/songs', views.SongsView.as_view(), name='songs'),
    path('api/songs/<int:song_id>', views.SongInfoView.as_view(), name='song_info'),
    path('api/songs/search/', views.SongSearchView.as_view(), name='song_search'),
    path('api/artists', views.ArtistsView.as_view(), name='artists'),
    path('api/artists/<int:artist_id>', views.ArtistInfoView.as_view(), name='artist_info'),
    path('api/albums', views.AlbumsView.as_view(), name='albums'),
    path('api/albums/<int:album_id>', views.AlbumInfoView.as_view(), name='album_info'),
]


from datetime import timedelta, datetime
from typing import Optional
from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from sqlalchemy.orm import Session
from jose import jwt, JWTError

from app import database, models
from app.schemas import TokenData
from app.config import settings

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")

SECRET_KEY = settings.secret_key
ALGORITHM = settings.algorithm
ACCESS_TOKEN_EXPIRE_MINUTES = settings.access_token_expire_minutes

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

def verify_access_token(token: str, credentials_exception):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        user_id: str = payload.get("user_id")
        if user_id is None:
            raise credentials_exception
        token_data = TokenData(id=user_id)
    except JWTError:
        raise credentials_exception
    return token_data

def get_current_user(
    token: str = Depends(oauth2_scheme), db: Session = Depends(database.get_db)
):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    token = verify_access_token(token, credentials_exception)
    user = db.query(models.User).filter(models.User.id == token.id).first()
    return user


from fastapi import APIRouter, Body, Depends, Request, status
from fastapi.responses import JSONResponse
from jarvis.db.database import DataBase, get_database
from jarvis.core import config, utils
from jarvis.lib import TwilioHelper
from typing import Dict
from twilio.rest import Client
import jarvis.crud as crud
import jarvis.models as model
import jarvis.core.text_responses as text

router = APIRouter()
twilio_helper = TwilioHelper()
client = Client(config.TWILIO_ACCOUNT_SID, config.TWILIO_ACCOUNT_AUTH_TOKEN)

@router.post("/add")
async def add_item_to_cart(request: Request, db: DataBase = Depends(get_database)):
    async with db.pool.acquire() as conn:
        body = await request.form()
        parsed_body = dict(body)
        cart_item = model.CartItem(**parsed_body)
        normalized_cart_item = await utils.normalize_cart_item_model(conn, cart_item)
        cart_item_name = normalized_cart_item.get("name")
        item_quantity = normalized_cart_item.get("quantity")

        success_message = text.add_item_success(cart_item_name, item_quantity)
        shopping_cart_message = text.shopping_cart_info(1)
        msg = "".join([success_message, shopping_cart_message])
        return twilio_helper.compose_mesage(msg)
        # Make potentially a new helper class that has add item
        # because you have to then convert this to a message after etc
        # shopping_cart = model.ShoppingCart(**payload)

        # return None

@router.get("/menu/{item_type}")
async def get_menu(
    item_type: str, db: DataBase = Depends(get_database),
):
    async with db.pool.acquire() as conn:
        try:
            items = await crud.get_all_item_by_type(conn, item_type)
            message_list = [utils.item_model_to_message(item) for item in items]
            message = "\n".join(message_list)
            twilio_message = twilio_helper.compose_mesage(message)
            return twilio_message

        except UserWarning as warning:
            return JSONResponse(
                status_code=status.HTTP_202_ACCEPTED, content=str(warning)
            )

@router.post("/checkout")
async def checkout_cart(
    payload: Dict = Body(...), db: DataBase = Depends(get_database)
):
    pass

@router.post("/sms")
async def get_twilio_text():
    resp = ":)"
    return utils.create_text_response(resp)

@router.get("/test")
async def twilio_test(payload: Dict = Body(...)):

    message = client.messages.create(
        body="Jarvis test",
        messaging_service_sid=config.TWILIO_ACCOUNT_MESSAGING_SID,
        to=config.TO_PHONE_NUMBER,
    )

    return message.sid


#!/usr/bin/env python
# Code property of Matteo Scanavino - matteo.svanavino@gmail.it
# Minor changes by Iris David Du Mutel
import rospy
# from std_msgs.msg import Float32MultiArray
from myrobot.msg import vect_msg
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
import cv2
# import cv2.cv
import os
import math
import numpy as np
#import pyrealsense2 as rs
import message_filters
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge, CvBridgeError
import imutils #collection of OpenCV and Python convenience functions
from collections import deque
from scipy.spatial.transform import Rotation as R

def green_ball():
    rospy.init_node('realsense_behaviour', anonymous=True)
    pub = rospy.Publisher('gb_vect', vect_msg, queue_size=10)
    color_sub = message_filters.Subscriber('camera/color/image_raw',Image)
    # depth_sub = message_filters.Subscriber('camera/depth/image_raw',Image)
    x_sub = message_filters.Subscriber('/odom',Odometry)
    ts = message_filters.ApproximateTimeSynchronizer([color_sub, x_sub], queue_size=10,slop=0.1)
    ts.registerCallback(callback,pub)
    rospy.spin()

def callback(color_raw, x_sub,pub):
    vect = [0, 0]
    msg = vect_msg()
    bridge = CvBridge()
    greenLower = (29, 86, 6)
    greenUpper = (64, 255, 255)
    # realsense min and max distance
    try:
            color_image = bridge.imgmsg_to_cv2(color_raw, "bgr8")
    except CvBridgeError as e:
            print(e)
    Xest = x_sub
    # # Variable assignation:

    [yaw, pitch, roll] = get_rotation(Xest)
    psi_est = yaw*180/math.pi
    frame = imutils.resize(color_image, width=600)
    blurred = cv2.GaussianBlur(frame, (11, 11), 0)
    hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
    # construct a mask for the color "green", then perform
	# a series of dilations and erosions to remove any small
    # blobs left in the mask
    mask = cv2.inRange(hsv, greenLower, greenUpper)
    mask = cv2.erode(mask, None, iterations=2)
    mask = cv2.dilate(mask, None, iterations=2)
    # find contours in the mask and initialize the current
    # (x, y) center of the ball
    cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    center = None
    if len(cnts) > 0:
        # find the largest contour in the mask, then use
        # it to compute the minimum enclosing circle and
        # centroid
        c = max(cnts, key=cv2.contourArea)
        ((x, y), radius) = cv2.minEnclosingCircle(c)
        if x<280:
            vect[0]=90
            vect[1]=0
        elif x>305:
            vect[0]=-90
            vect[1]=0
        else:
            if radius<100:
                vect[0]=psi_est
                vect[1]=0.8
            else:
                vect[0]= psi_est
                vect[1]=0
        M = cv2.moments(c)
        center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
        print(center)
        # only proceed if the radius meets a minimum size
        if radius > 10:
            # draw the circle and centroid on the frame,
            # then update the list of tracked points
            cv2.circle(frame, (int(x), int(y)), int(radius),
                (0, 255, 255), 2)
            print('radius=', radius)
            cv2.circle(frame, center, 5, (0, 0, 255), -1)
    else:
        print('out of frame')
    # show the frame to our screen
    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF
    # Send data
    msg.header.stamp = rospy.Time.now()
    msg.angle = vect[0]
    msg.value = vect[1]
    rospy.loginfo('Realsense vector data sent')
    pub.publish(msg)

def get_rotation(Xest):
    orientation_q = Xest.pose.pose.orientation
    orientation_list = [orientation_q.x, orientation_q.y, orientation_q.z, orientation_q.w]
    r = R.from_quat(orientation_list)
    EuAn = r.as_euler('zyx', degrees=False)
    return EuAn
if __name__ == '__main__':
    try:
        green_ball()
    except rospy.ROSInterruptException:
        pass


from database import db
from flask import request
from middleware.auth import login_required, admin_only
from models.guild import Guild
from typing import Dict, Optional, Tuple

def check_request(req: request, id_only: Optional[bool] = False) -> int | Tuple[int, str, bool] | Tuple[Dict[str, str], int]:
    # Check request body
    guild_name = ''
    guild_manage_threads = False
    try:
        guild_id = req.json['id']

        if not isinstance(guild_id, int):
            raise ValueError('id must be an integer')

        if not id_only:
            guild_name = req.json.get('name', guild_name)
            guild_manage_threads = req.json.get('manage_threads', guild_manage_threads)

            if 'name' in req.json and not isinstance(guild_name, str):
                raise ValueError('name must be a string')
            if 'name' in req.json and not 0 < len(guild_name) < 256:
                raise ValueError('name must be between 0 to 256 characters long')
            if 'manage_threads' in req.json and not isinstance(guild_manage_threads, bool):
                raise ValueError('manage_threads must be a boolean')
    except KeyError as e:
        return {
            'success': False,
            'error': f'Missing key in request body: {e}'
        }, 400
    except ValueError as e:
        return {
            'success': False,
            'error': f'Bad value: {e}'
        }, 400
    else:
        if id_only:
            return guild_id
        else:
            return guild_id, guild_name, guild_manage_threads

@admin_only
def add_guild():
    # Check request body
    check_result = check_request(request)
    if isinstance(check_result[0], dict):
        return check_result
    guild_id, guild_name, guild_manage_threads = check_result

    # Check if guild is already in DB
    guild = Guild.query.get(guild_id)
    if guild is not None:
        return {
            'success': False,
            'error': 'Guild already exists'
        }, 409

    # Create guild
    guild = Guild(
        id=guild_id,
        name=guild_name,
        manage_threads=guild_manage_threads
    )

    # Add to DB
    db.session.add(guild)
    db.session.commit()
    return {
        'success': True
    }

@admin_only
def update_guild():
    # Check request body
    check_result = check_request(request)
    if isinstance(check_result[0], dict):
        return check_result
    guild_id, guild_name, guild_manage_threads = check_result

    # Check if guild is already in DB
    guild = Guild.query.get(guild_id)
    if guild is None:
        return {
            'success': False,
            'error': f'Guild {guild_id} does not exist'
        }, 404

    # Update existing guild
    if 'name' in request.json:
        guild.name = guild_name
    if 'manage_threads' in request.json:
        guild.manage_threads = guild_manage_threads

    # Commit
    db.session.commit()
    return {
        'success': True
    }, 200

@admin_only
def delete_guild():
    # Check request body
    check_result = check_request(request, id_only=True)
    if isinstance(check_result, tuple) and isinstance(check_result[0], dict):
        return check_result
    guild_id = check_result

    # Check if guild is in DB
    guild = Guild.query.get(guild_id)
    if guild is not None:
        # Delete user
        db.session.delete(guild)
        db.session.commit()

    return {
        'success': True
    }, 200

@login_required
def get_guild():
    # Check request body
    check_result = check_request(request, id_only=True)
    if isinstance(check_result, tuple) and isinstance(check_result[0], dict):
        return check_result
    guild_id = check_result

    # Check if guild is in DB
    guild = Guild.query.get(guild_id)
    if guild is None:
        return {
            'success': False,
            'error': 'Guild not found'
        }, 404

    # Return guild data
    return {
        'success': True,
        'guild': {
            'name': guild.name,
            'manage_threads': guild.manage_threads
        }
    }, 200


from __future__ import annotations

from textwrap import dedent

import pytest

from pants.backend.java.target_types import JavaSourcesGeneratorTarget
from pants.backend.java.target_types import rules as target_types_rules
from pants.core.util_rules import config_files, source_files
from pants.engine.addresses import Address, Addresses
from pants.jvm.resolve.common import Coordinate
from pants.jvm.resolve.coursier_fetch import NoCompatibleResolve
from pants.jvm.resolve.coursier_fetch import rules as coursier_fetch_rules
from pants.jvm.resolve.key import CoursierResolveKey
from pants.jvm.target_types import DeployJarTarget, JvmArtifactTarget
from pants.jvm.testutil import maybe_skip_jdk_test
from pants.jvm.util_rules import rules as util_rules
from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, QueryRule, RuleRunner, engine_error

NAMED_RESOLVE_OPTIONS = (
    '--jvm-resolves={"one": "coursier_resolve.lockfile", "two": "coursier_resolve.lockfile"}'
)
DEFAULT_RESOLVE_OPTION = "--jvm-default-resolve=one"

@pytest.fixture
def rule_runner() -> RuleRunner:
    rule_runner = RuleRunner(
        rules=[
            *config_files.rules(),
            *coursier_fetch_rules(),
            *source_files.rules(),
            *util_rules(),
            *target_types_rules(),
            QueryRule(CoursierResolveKey, (Addresses,)),
        ],
        target_types=[DeployJarTarget, JavaSourcesGeneratorTarget, JvmArtifactTarget],
    )
    rule_runner.set_options(
        args=[
            NAMED_RESOLVE_OPTIONS,
            DEFAULT_RESOLVE_OPTION,
        ],
        env_inherit=PYTHON_BOOTSTRAP_ENV,
    )
    return rule_runner

def assert_resolve(
    expected_resolve: str,
    rule_runner: RuleRunner,
    root_one_resolve: str,
    root_two_resolve: str,
    leaf_resolve: str,
) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                f"""\
                deploy_jar(name='root_one', main='Ex', dependencies=[':leaf'], resolve='{root_one_resolve}')
                deploy_jar(name='root_two', main='Ex', dependencies=[':leaf'], resolve='{root_two_resolve}')
                jvm_artifact(
                  name='leaf',
                  group='ex',
                  artifact='ex',
                  version='0.0.0',
                  resolve='{leaf_resolve}',
                )
                """
            ),
            "coursier_resolve.lockfile": "[]",
        }
    )
    resolve_key = rule_runner.request(
        CoursierResolveKey,
        # NB: Although it will not happen for `deploy_jars` in production, we resolve two of them
        # together here to validate the handling of multiple roots, which _can_ happen for things
        # like the `repl` goal, and other goals which create an adhoc merged Classpath.
        [
            Addresses(
                [
                    Address(spec_path="", target_name="root_one"),
                    Address(spec_path="", target_name="root_two"),
                ]
            )
        ],
    )
    assert resolve_key.name == expected_resolve

@maybe_skip_jdk_test
def test_all_matching(rule_runner: RuleRunner) -> None:
    assert_resolve("one", rule_runner, "one", "one", "one")

@maybe_skip_jdk_test
def test_no_matching_for_root(rule_runner: RuleRunner) -> None:
    with engine_error(NoCompatibleResolve):
        assert_resolve("n/a", rule_runner, "one", "two", "two")

@maybe_skip_jdk_test
def test_no_matching_for_leaf(rule_runner: RuleRunner) -> None:
    with engine_error(NoCompatibleResolve):
        assert_resolve("n/a", rule_runner, "one", "one", "two")

@pytest.mark.parametrize(
    "coord_str,expected",
    (
        ("group:artifact:version", Coordinate("group", "artifact", "version")),
        (
            "group:artifact:packaging:version",
            Coordinate("group", "artifact", "version", "packaging"),
        ),
        (
            "group:artifact:packaging:classifier:version",
            Coordinate("group", "artifact", "version", "packaging", "classifier"),
        ),
    ),
)
def test_from_coord_str(coord_str: str, expected: Coordinate) -> None:
    assert Coordinate.from_coord_str(coord_str) == expected


import torch
import csv
import pytorch_lightning as pl
from sys import platform
if platform == "linux":
    from pypesq import pesq
from pystoi import stoi
from math import isnan
from numpy import random

def check_inf_neginf_nan(tensor, error_msg):
    assert not torch.any(torch.isinf(tensor)), error_msg
    if tensor.dtype == torch.complex32 or tensor.dtype == torch.complex64 or tensor.dtype == torch.complex128:
        assert not torch.any(torch.isneginf(tensor.real)), error_msg
        assert not torch.any(torch.isneginf(tensor.imag)), error_msg
    else:
        assert not torch.any(torch.isneginf(tensor)), error_msg
    assert not torch.any(torch.isnan(tensor)), error_msg

def l2_norm(s1, s2):
    norm = torch.sum(s1*s2, -1, keepdim=True)

    return norm

# source https://arxiv.org/pdf/2008.00264.pdf
class SiSNR(object):
    def __call__(self, clean, estimate, eps=1e-8):
        dot = l2_norm(estimate, clean)
        norm = l2_norm(clean, clean)

        s_target =  (dot * clean)/(norm+eps)
        e_nosie = estimate - s_target

        target_norm = l2_norm(s_target, s_target)
        noise_norm = l2_norm(e_nosie, e_nosie)
        snr = 10*torch.log10((target_norm)/(noise_norm+eps)+eps)

        return torch.mean(snr)

# source https://github.com/chanil1218/DCUnet.pytorch/blob/2dcdd30804be47a866fde6435cbb7e2f81585213/train.py
class wSDR(object):
    def __call__(self, mixed, clean, clean_est, eps=2e-8):
        bsum = lambda x: torch.sum(x, dim=1)
        def mSDRLoss(orig, est):
            correlation = bsum(orig * est)
            energies = torch.norm(orig, p=2, dim=1) * torch.norm(est, p=2, dim=1)
            return -(correlation / (energies + eps))

        noise = mixed - clean
        noise_est = mixed - clean_est

        a = bsum(clean**2) / (bsum(clean**2) + bsum(noise**2) + eps)
        target_wSDR = a * mSDRLoss(clean, clean_est)
        noise_wSDR = (1 - a) * mSDRLoss(noise, noise_est)
        wSDR = target_wSDR + noise_wSDR
        return torch.mean(wSDR)

def cRM(S, Y, eps=1e-8):
    M_r_numer = (Y.real * S.real) + (Y.imag * S.imag)
    M_r_denom = torch.square(Y.real) + torch.square(Y.imag)
    M_r = M_r_numer / (M_r_denom + eps)

    M_i_numer = (Y.real * S.imag) - (Y.imag * S.real)
    M_i_denom = torch.square(Y.real) + torch.square(Y.imag)
    M_i = M_i_numer / (M_i_denom + eps)

    M = torch.complex(M_r, M_i)

    return M

def bound_cRM(cRM):
    target_noise_mask_mag = torch.abs(cRM)
    target_noise_mask_mag_tanh = torch.tanh(target_noise_mask_mag)
    target_noise_mag_tanh_real = target_noise_mask_mag_tanh * torch.cos(torch.angle(cRM))
    target_noise_mag_tanh_imag = target_noise_mask_mag_tanh * torch.sin(torch.angle(cRM))
    target_noise_mask_phase = torch.atan2(target_noise_mag_tanh_imag, target_noise_mag_tanh_real)

    target_noise_mask_real = target_noise_mask_mag_tanh * torch.cos(target_noise_mask_phase)
    target_noise_mask_imag = target_noise_mask_mag_tanh * torch.sin(target_noise_mask_phase)

    return torch.complex(target_noise_mask_real, target_noise_mask_imag)

def complex_mat_mult(A, B):
    outp_real = (A.real * B.real) - (A.imag * B.imag)
    outp_imag = (A.real * B.imag) + (A.imag * B.real)

    Y = torch.complex(outp_real, outp_imag)

    return Y

def complex_lrelu(input):
    # return torch.nn.functional.leaky_relu(input.real) + 1j*torch.nn.functional.leaky_relu(input.imag)
    return torch.complex(torch.nn.functional.leaky_relu(input.real), torch.nn.functional.leaky_relu(input.imag))

def apply_complex(fr, fi, input):
    # return (fr(input.real)[0]-fi(input.imag)[0]) + 1j*(fr(input.imag)[0]+fi(input.real)[0])
    return torch.complex(fr(input.real)-fi(input.imag), (fr(input.imag)+fi(input.real)))

# source https://github.com/huyanxin/DeepComplexCRN/blob/bc6fd38b0af9e8feb716c81ff8fbacd7f71ad82f/complexnn.py
class ComplexLSTM(torch.nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, bidirectional, batch_first, projection_dim=None):
        super(ComplexLSTM, self).__init__()

        self.input_dim = input_size
        self.rnn_units = hidden_size
        self.real_lstm = torch.nn.LSTM(input_size=self.input_dim, hidden_size=self.rnn_units, num_layers=num_layers,
                        bidirectional=bidirectional, batch_first=batch_first)
        self.imag_lstm = torch.nn.LSTM(input_size=self.input_dim, hidden_size=self.rnn_units, num_layers=num_layers,
                        bidirectional=bidirectional, batch_first=batch_first)
        if bidirectional:
            bidirectional=2
        else:
            bidirectional=1
        if projection_dim is not None:
            self.projection_dim = projection_dim
            self.r_trans = torch.nn.Linear(self.rnn_units*bidirectional, self.projection_dim)
            self.i_trans = torch.nn.Linear(self.rnn_units*bidirectional, self.projection_dim)
        else:
            self.projection_dim = None

    def forward(self, inputs):
        if isinstance(inputs,list):
            real, imag = inputs.real, inputs.imag
        elif isinstance(inputs, torch.Tensor):
            real, imag = inputs.real, inputs.imag
        r2r_out = self.real_lstm(real)[0]
        r2i_out = self.imag_lstm(real)[0]
        i2r_out = self.real_lstm(imag)[0]
        i2i_out = self.imag_lstm(imag)[0]
        real_out = r2r_out - i2i_out
        imag_out = i2r_out + r2i_out
        if self.projection_dim is not None:
            real_out = self.r_trans(real_out)
            imag_out = self.i_trans(imag_out)
        return torch.complex(real_out, imag_out)

    def flatten_parameters(self):
        self.imag_lstm.flatten_parameters()
        self.real_lstm.flatten_parameters()

def mag_phase_2_wave(mag, phase, config):
    real = mag * torch.cos(phase)
    imag = mag * torch.sin(phase)
    comp = torch.complex(real, imag)
    comp = torch.nn.functional.pad(comp, (0,0,0,1))

    audio = torch.istft(comp, n_fft=config.fft_size, hop_length=config.hop_length, \
            win_length=config.window_length, normalized=config.normalise_stft)

    return audio

def calc_metric(clean_audio, predict_audio, config, metric):
    metric_arr = []
    for i in range(predict_audio.shape[0]):
        metric_i = metric(clean_audio[i,:].cpu().numpy(), predict_audio[i,:].cpu().numpy(), config.sr)
        if not isnan(metric_i):
            metric_arr.append(metric_i)
    pesq_av = float(sum(metric_arr)) / max(len(metric_arr), 1)

    return pesq_av

def calc_loss(self, target_noise_mask, predict_noise_mask, \
                predict_noise_audio, predict_clean_audio,
                noise_audio, noisy_audio, clean_audio):

    if self.hparams['noise_loss_type'] == 0:
        noise_loss_orig = self.config.L1(target_noise_mask, predict_noise_mask)
    elif self.hparams['noise_loss_type'] == 1:
        noise_loss_orig = self.config.wSDR(noisy_audio, noise_audio, predict_noise_audio)
    elif self.hparams['noise_loss_type'] == 2:
        noise_loss_orig = self.config.L1(target_noise_mask, predict_noise_mask) + \
            self.config.L1(noise_audio, predict_noise_audio)
    elif self.hparams['noise_loss_type'] == 3:
        noise_loss_orig = self.config.wSDR(noisy_audio, noise_audio, predict_noise_audio) + \
            self.config.L1(noise_audio, predict_noise_audio)
    elif self.hparams['noise_loss_type'] == 4:
        noise_loss_orig = self.config.wSDR(noisy_audio, noise_audio, predict_noise_audio) + \
            self.config.L1(target_noise_mask, predict_noise_mask)
    elif self.hparams['noise_loss_type'] == 5:
        if target_noise_mask.dtype == torch.complex32 or target_noise_mask.dtype == torch.complex64 or target_noise_mask.dtype == torch.complex128:
            noise_loss_orig = self.config.wSDR(noisy_audio, noise_audio, predict_noise_audio) + \
            self.config.mse(target_noise_mask.real, predict_noise_mask.real) + \
            self.config.mse(target_noise_mask.imag, predict_noise_mask.imag)
        else:
            noise_loss_orig = self.config.wSDR(noisy_audio, noise_audio, predict_noise_audio) + \
                self.config.mse(target_noise_mask, predict_noise_mask)
    noise_loss = (self.hparams['noise_alpha'] * noise_loss_orig)

    if self.hparams['speech_loss_type'] == 0:
        speech_loss_orig = -self.config.SiSNR(clean_audio, predict_clean_audio)
    elif self.hparams['speech_loss_type'] == 1:
        speech_loss_orig_small = torch.mean(self.config.CDPAM.forward(clean_audio, predict_clean_audio))
        speech_loss_orig = speech_loss_orig_small * 10e5
    speech_loss = (self.hparams['speech_alpha'] * speech_loss_orig)

    total_loss = noise_loss + speech_loss

    return noise_loss, speech_loss, total_loss

def train_batch_2_loss(self, train_batch, batch_idx, dtype):
    noise_data, noisy_data, clean_data, id = train_batch

    check_inf_neginf_nan(clean_data, "Found inf, neginf or nan in clean data STFT!")
    check_inf_neginf_nan(noise_data, "Found inf, neginf or nan in noise data STFT!")
    check_inf_neginf_nan(noisy_data, "Found inf, neginf or nan in noisy data STFT!")

    noise_mag = torch.abs(noise_data)
    noise_phase = torch.angle(noise_data)
    noisy_mag = torch.abs(noisy_data)
    noisy_phase = torch.angle(noisy_data)
    clean_mag = torch.abs(clean_data)
    clean_phase = torch.angle(clean_data)
    noise_audio = mag_phase_2_wave(noise_mag, noise_phase, self.config)
    noisy_audio = mag_phase_2_wave(noisy_mag, noisy_phase, self.config)
    clean_audio = mag_phase_2_wave(clean_mag, clean_phase, self.config)

    if dtype == "real":
        target_noise_mask = torch.sigmoid(noise_mag / noisy_mag)

        noisy_mag_scaled = (noisy_mag - self.config.data_minR) / (self.config.data_maxR - self.config.data_minR)
        predict_noise_mask = self(noisy_mag_scaled)

        predict_noise_mag = noisy_mag * predict_noise_mask
        predict_clean_mag = noisy_mag - predict_noise_mag
        predict_noise_audio = mag_phase_2_wave(predict_noise_mag, noisy_phase, self.config)
        predict_clean_audio = mag_phase_2_wave(predict_clean_mag, noisy_phase, self.config)

    elif dtype == "complex":
        target_noise_mask_out = cRM(noise_data, noisy_data)
        target_noise_mask = bound_cRM(target_noise_mask_out)

        # noisy_data_standardised = (noisy_data - torch.mean(noisy_data)) / torch.std(noisy_data)
        noisy_data_scaled = torch.view_as_complex((2 * ((torch.view_as_real(noisy_data) - self.config.data_minC) /
                    (self.config.data_maxC - self.config.data_minC))) - 1)
        predict_noise_mask_out = self(noisy_data_scaled)
        predict_noise_mask = bound_cRM(predict_noise_mask_out)

        predict_noise_data = complex_mat_mult(noisy_data, predict_noise_mask)
        predict_clean_data = noisy_data - predict_noise_data
        predict_noise_audio = mag_phase_2_wave(torch.abs(predict_noise_data), \
                                torch.angle(predict_noise_data), self.config)
        predict_clean_audio = mag_phase_2_wave(torch.abs(predict_clean_data), \
                                torch.angle(predict_clean_data), self.config)

    noise_loss, speech_loss, train_loss = calc_loss(self,
                                                    target_noise_mask=target_noise_mask,
                                                    predict_noise_mask=predict_noise_mask,
                                                    predict_noise_audio=predict_noise_audio,
                                                    predict_clean_audio=predict_clean_audio,
                                                    noise_audio=noise_audio,
                                                    noisy_audio=noisy_audio,
                                                    clean_audio=clean_audio)

    return noise_loss, speech_loss, train_loss

def val_batch_2_metric_loss(self, val_batch, val_idx, dtype):
    noise_data, noisy_data, clean_data, id = val_batch

    check_inf_neginf_nan(clean_data, "Found inf, neginf or nan in clean data STFT!")
    check_inf_neginf_nan(noise_data, "Found inf, neginf or nan in noise data STFT!")
    check_inf_neginf_nan(noisy_data, "Found inf, neginf or nan in noisy data STFT!")

    noise_mag = torch.abs(noise_data)
    noise_phase = torch.angle(noise_data)
    noisy_mag = torch.abs(noisy_data)
    noisy_phase = torch.angle(noisy_data)
    clean_mag = torch.abs(clean_data)
    clean_phase = torch.angle(clean_data)
    noise_audio = mag_phase_2_wave(noise_mag, noise_phase, self.config)
    noisy_audio = mag_phase_2_wave(noisy_mag, noisy_phase, self.config)
    clean_audio = mag_phase_2_wave(clean_mag, clean_phase, self.config)

    if dtype == "real":
        target_noise_mask = torch.sigmoid(noise_mag / noisy_mag)

        noisy_mag_scaled = (noisy_mag - self.config.data_minR) / (self.config.data_maxR - self.config.data_minR)
        predict_noise_mask = self(noisy_mag_scaled)

        predict_noise_mag = noisy_mag * predict_noise_mask
        predict_clean_mag = noisy_mag - predict_noise_mag

        predict_clean_audio = mag_phase_2_wave(predict_clean_mag, noisy_phase, self.config)
        predict_noise_audio = mag_phase_2_wave(predict_noise_mag, noisy_phase, self.config)

    elif dtype == "complex":
        target_noise_mask_out = cRM(noise_data, noisy_data)
        target_noise_mask = bound_cRM(target_noise_mask_out)

        # noisy_data_standardised = (noisy_data - torch.mean(noisy_data)) / torch.std(noisy_data)
        noisy_data_scaled = torch.view_as_complex((2 * ((torch.view_as_real(noisy_data) - self.config.data_minC) /
                    (self.config.data_maxC - self.config.data_minC))) - 1)
        predict_noise_mask_out = self(noisy_data_scaled)
        predict_noise_mask = bound_cRM(predict_noise_mask_out)

        predict_noise_data = complex_mat_mult(noisy_data, predict_noise_mask)
        predict_clean_data = noisy_data - predict_noise_data
        predict_clean_audio = mag_phase_2_wave(torch.abs(predict_clean_data), \
                                torch.angle(predict_clean_data), self.config)
        predict_noise_audio = mag_phase_2_wave(torch.abs(predict_noise_data), \
                                torch.angle(predict_noise_data), self.config)

    if platform == "linux":
        pesq_av = calc_metric(clean_audio, predict_clean_audio, self.config, pesq)
    else:
        pesq_av = 1
    stoi_av = calc_metric(clean_audio, predict_clean_audio, self.config, stoi)

    noise_loss, speech_loss, val_loss = calc_loss(self,
                                                    target_noise_mask=target_noise_mask,
                                                    predict_noise_mask=predict_noise_mask,
                                                    predict_noise_audio=predict_noise_audio,
                                                    predict_clean_audio=predict_clean_audio,
                                                    noise_audio=noise_audio,
                                                    noisy_audio=noisy_audio,
                                                    clean_audio=clean_audio)

    return noise_loss, speech_loss, val_loss, pesq_av, stoi_av, \
                    predict_noise_audio, predict_clean_audio, \
                    noise_audio, noisy_audio, clean_audio

def test_batch_2_metric_loss(self, test_batch, test_idx, dtype):
    noise_data, noisy_data, clean_data, id, start_point = test_batch

    noise_mag = torch.abs(noise_data)
    noise_phase = torch.angle(noise_data)
    noisy_mag = torch.abs(noisy_data)
    noisy_phase = torch.angle(noisy_data)
    clean_mag = torch.abs(clean_data)
    clean_phase = torch.angle(clean_data)
    noise_audio = mag_phase_2_wave(noise_mag, noise_phase, self.config)
    noisy_audio = mag_phase_2_wave(noisy_mag, noisy_phase, self.config)
    clean_audio = mag_phase_2_wave(clean_mag, clean_phase, self.config)

    if dtype == "real":
        target_noise_mask = torch.sigmoid(noise_mag / noisy_mag)

        noisy_mag_scaled = (noisy_mag - self.config.data_minR) / (self.config.data_maxR - self.config.data_minR)
        predict_noise_mask = self(noisy_mag_scaled)

        predict_noise_mag = noisy_mag * predict_noise_mask
        predict_clean_mag = noisy_mag - predict_noise_mag

        predict_clean_audio = mag_phase_2_wave(predict_clean_mag, noisy_phase, self.config)
        predict_noise_audio = mag_phase_2_wave(predict_noise_mag, noisy_phase, self.config)

    elif dtype == "complex":
        target_noise_mask_out = cRM(noise_data, noisy_data)
        target_noise_mask = bound_cRM(target_noise_mask_out)

        # noisy_data_standardised = (noisy_data - torch.mean(noisy_data)) / torch.std(noisy_data)
        noisy_data_scaled = torch.view_as_complex((2 * ((torch.view_as_real(noisy_data) - self.config.data_minC) /
                    (self.config.data_maxC - self.config.data_minC))) - 1)
        predict_noise_mask_out = self(noisy_data_scaled)
        predict_noise_mask = bound_cRM(predict_noise_mask_out)

        predict_noise_data = complex_mat_mult(noisy_data, predict_noise_mask)
        predict_clean_data = noisy_data - predict_noise_data
        predict_clean_audio = mag_phase_2_wave(torch.abs(predict_clean_data), \
                                torch.angle(predict_clean_data), self.config)
        predict_noise_audio = mag_phase_2_wave(torch.abs(predict_noise_data), \
                                torch.angle(predict_noise_data), self.config)

    noise_audio = mag_phase_2_wave(noise_mag, noise_phase, self.config)
    noisy_audio = mag_phase_2_wave(noisy_mag, noisy_phase, self.config)

    if platform == "linux":
        pesq_av = calc_metric(clean_audio, predict_clean_audio, self.config, pesq)
    else:
        pesq_av = 1
    stoi_av = calc_metric(clean_audio, predict_clean_audio, self.config, stoi)

    noise_loss, speech_loss, test_loss = calc_loss(self,
                                                    target_noise_mask=target_noise_mask,
                                                    predict_noise_mask=predict_noise_mask,
                                                    predict_noise_audio=predict_noise_audio,
                                                    predict_clean_audio=predict_clean_audio,
                                                    noise_audio=noise_audio,
                                                    noisy_audio=noisy_audio,
                                                    clean_audio=clean_audio)

    return noise_loss, speech_loss, test_loss, pesq_av, stoi_av, \
                    predict_noise_audio, predict_clean_audio, \
                    noise_audio, noisy_audio, clean_audio, id, start_point

def epoch_end(self, outputs, type):

    no_of_batches = len(outputs)
    random_batches = random.choice(no_of_batches, size=min(self.config.val_log_sample_size, no_of_batches), replace=False)

    no_of_samples = min(self.config.data_params['batch_size'],
                    outputs[-1]['clean'].shape[0],
                    outputs[-1]['predict_clean'].shape[0],
                    outputs[-1]['noise'].shape[0],
                    outputs[-1]['predict_noise'].shape[0],
                    outputs[-1]['noisy'].shape[0])
    random_samples = random.choice(no_of_samples, size=min(self.config.val_log_sample_size, no_of_samples), replace=False)

    for i, ridx in enumerate(range(min(self.config.val_log_sample_size, no_of_samples))):
        clean_sample = outputs[random_batches[ridx]]['clean'][random_samples[ridx],:]
        predict_clean_sample = outputs[random_batches[ridx]]['predict_clean'][random_samples[ridx],:]
        noise_sample = outputs[random_batches[ridx]]['noise'][random_samples[ridx],:]
        predict_noise_sample = outputs[random_batches[ridx]]['predict_noise'][random_samples[ridx],:]
        noisy_sample = outputs[random_batches[ridx]]['noisy'][random_samples[ridx],:]

        self.logger.experiment.add_audio("clean({})/{}".format(type, i),
                                        clean_sample,
                                        self.global_step,
                                        sample_rate=self.config.sr)
        self.logger.experiment.add_audio("predict_clean({})/{}".format(type, i),
                                        predict_clean_sample,
                                        self.global_step,
                                        sample_rate=self.config.sr)
        self.logger.experiment.add_audio("noise({})/{}".format(type, i),
                                        noise_sample,
                                        self.global_step,
                                        sample_rate=self.config.sr)
        self.logger.experiment.add_audio("predict_noise({})/{}".format(type, i),
                                        predict_noise_sample,
                                        self.global_step,
                                        sample_rate=self.config.sr)
        self.logger.experiment.add_audio("noisy({})/{}".format(type, i),
                                        noisy_sample,
                                        self.global_step,
                                        sample_rate=self.config.sr)

class InputMonitor(pl.Callback):

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        if (batch_idx + 1) % trainer.log_every_n_steps == 0:
            noise_real = batch[0].real
            noise_imag = batch[0].imag
            noisy_real = batch[1].real
            noisy_imag = batch[1].imag
            clean_real = batch[2].real
            clean_imag = batch[2].imag
            logger = trainer.logger
            logger.experiment.add_histogram("noise data real", noise_real, global_step=trainer.global_step)
            logger.experiment.add_histogram("noise data imag", noise_imag, global_step=trainer.global_step)
            logger.experiment.add_histogram("noisy data real", noisy_real, global_step=trainer.global_step)
            logger.experiment.add_histogram("noisy data imag", noisy_imag, global_step=trainer.global_step)
            logger.experiment.add_histogram("clean data real", clean_real, global_step=trainer.global_step)
            logger.experiment.add_histogram("clean data imag", clean_imag, global_step=trainer.global_step)

class CheckBatchGradient(pl.Callback):
    def on_train_start(self, trainer, model):
        n = 0

        example_input = model.example_input_array.to(model.device)
        example_input.requires_grad = True

        model.zero_grad()
        output = model(example_input)
        output[n].abs().sum().backward()

        zero_grad_inds = list(range(example_input.size(0)))
        zero_grad_inds.pop(n)

        if example_input.grad[zero_grad_inds].abs().sum().item() > 0:
            raise RuntimeError("Your model mixes data across the batch dimension!")


import requests
from bs4 import BeautifulSoup
import pandas as pd
from os import listdir, remove
import datetime as dt
from time import sleep

from MainMethods import getInfo, showDays
from conf import INT, INF, URL, LOC, NINF, LOC2,\
    chosenF, errorsF, doneF

"""
The information for saved days is checked
and old files deleted.
"""
oldflags = [f for f in listdir(LOC2) if f[0]== "F"]
if oldflags:
    for f in oldflags:
        remove(f"{LOC2}\{f}")

saved= listdir(LOC)
if saved:
    saved = [f.split(".csv")[0] for f in saved]
    ints = [int(f.split("y")[-1]) for f in saved]
    for i, f in enumerate(saved):
        if ints[i] < dt.datetime.today().day and max(ints) -ints[i] <9:
            remove(f"{LOC}\{f}.csv")
            saved.remove(f)

def flagIt():
    now = dt.datetime.timestamp(dt.datetime.now())
    name = fr"{LOC2}\F{str(int(now))}.txt"
    with open(name, "w") as F:
        pass
    sleep(2)
    return name

def unflagIt(name):
    remove(name)

def checkWait():
    flag = [f for f in listdir(LOC2) if f[0] == "T"]
    if flag:
        while flag[0] in listdir(LOC2):
            print("Wait, Sign Up in process....")
            sleep(5)

checkWait()

flag = flagIt()
Chosen = pd.read_csv(chosenF)
Done = pd.read_csv(doneF)[NINF[-2]].to_list()
if Done:
    print(f"These SignUps are done and should be confirmed by email:\n"
          f"{Chosen[Chosen[NINF[-2]].isin(Done)][[NINF[0], NINF[1], NINF[2]]].to_string(index=False)}\n\n"
          f"-------------------------------------------")
    Chosen.drop(Chosen[Chosen[NINF[-2]].isin(Done)].index, inplace=True)
    pd.DataFrame(columns= NINF).to_csv(doneF, index= False)

Errors = pd.read_csv(errorsF)[NINF[-2]].to_list()
if Errors:
    print(f"The sign up for these classes failed:\n"
          f"{Errors.iloc[:,:3]}\n"
          f"Please check manually if you are still interested and "
          f"allow them to be deleted from the program.")
    conf = "n"
    while conf.lower() != "y":
        conf = input("Allow? (y/n):")
        if conf.lower() == "y":
            Errors = pd.DataFrame(columns= NINF)
            Errors.to_csv(errorsF, index = False)
            Chosen.drop(Chosen[Chosen[NINF[-2]].isin(Errors)].index, inplace=True)
        else:
            conf = input("There is no benefit in keeping them !\n"
                         "Are you sure "
                         "you don't want to let them go?\n"
                         "(y/n):")

Chosen.to_csv(chosenF, index= False)

"""
This uses requests and beautiful soup to setup
the iterators.
"""
r = requests.get(URL)
soup = BeautifulSoup(r.text, "lxml")
classes = soup.find(id= "classes")
days = classes.find_all(class_= "scheduleDay")[:8]

"""
The following loop gets the basic info from the websites
and keeps it in the dictionary DFs as DataFrames
"""
DFs = {}
for day in days:
    date= day["id"]
    if date in saved:
        continue
    DFs[date] = pd.DataFrame(columns= INF)

    # iterate over each class in the day
    dayclss = day.find_all("div")
    for clss in dayclss:
        #then within each class I select the link in "schedSignup"
        if any(x in clss["class"] for x in INT):
            link = clss.find(class_= "schedSignup").a["href"]
            inf = getInfo(link)
            DFs[date] = DFs[date].append(pd.Series(inf, index= INF), ignore_index=True)

"""
This condition runs the showDays loop to check
each new day's classes for availability and presents the options
"""
num = 0
NewDF = pd.DataFrame(columns= NINF)
if DFs:
    result = showDays(DFs, num, NewDF)
    NewDF, num  = result[0], result[1]

#############

"""
Here, the requests waiting in the 'chosen' csv file
are presented and offered for cancellation
"""

# this just sets up the sig and UId variables for
sigs= [f for f in listdir(LOC2) if f[:3]== "Sig"]
if sigs:
    with open(f"{LOC2}\{sigs[0]}", "r") as s:
        UId = int(s.read())

                        ##### Cancel
if Chosen.shape[0]:
    print(f"\n============== OPTIONS TO CANCEL ======================\n"
          f"These are signups that are waiting to be executed:\n\n"
          f"{Chosen.iloc[:,:3]}\n\n"
          f"Type in the row number on the left if you want to cancel it, seperate with commas\n"
          f"Otherwise, just hit enter and confirm\n")

    confirm = "n"
    while confirm.lower() != "y":
        inp = input("CANCEL:")
        if inp:
            try:
                inp = list(map(int, inp.split(",")))
                print(f"cancel these:\n"
                      f"{Chosen.loc[inp, [NINF[0], NINF[1], NINF[2]]]}")
                confirm = input("Confirm (y/n):")
                if confirm.lower() == "y":
                    Chosen.drop(inp, inplace=True)
            except:
                print(f"There seems to be a mistake in your input,\n"
                      f"please don't type any unnecessary commas, spaces or words.")
        else:
            confirm = input("Keep all (y/n):")

"""
If there are newly available classes:
the following while loop will get requests and
add the newly chosen ones to the 'chosen' csv file

It will also give Unique IDs to each class based on
the UId variable retrieved from the Signal File (SigA or SigB)
"""

                        ##### Choose
if num:
    print(f"=====================================\n"
          f"The column on the RIGHT of each list contains the code to choose the class\n"
          f"please type in your choice(s)"
          f"(seperate codes with commas if you want multiple, hit enter if you want none.)\n")
    confirm = "n"
    while confirm.lower() != "y":
        choice = input("Choice:")
        if choice:
            try:
                choice = list(map(int,choice.split(",")))
                chosen = NewDF[NewDF[NINF[-1]].isin(choice)].copy()
                if max(choice) <= NewDF[NINF[-1]].max():
                    print(f"These are your new choices:\n"
                          f"{chosen.iloc[:,:3].to_string(index= False)}\n")

                    if Chosen.shape[0]:
                        print(f"These are still waiting to be executed:\n"
                              f"{Chosen.iloc[:, :3].to_string(index=False)}\n")
                    else:
                        print(f"There are no signups waiting.")
                    confirm = input("Confirm (y/n):")

                else:
                    print(f"You may have forgotten a comma or got the wrong number,\n"
                          f"please try again")

            except:
                print(f"There seems to be a mistake in your input,\n"
                      f"please don't type any unnecessary commas, spaces or words.")

        else:
            print(f"You chose none.")
            chosen = pd.DataFrame()
            if Chosen.shape[0]:
                print(f"These are still waiting to be executed:\n"
                      f"{Chosen.iloc[:, :3].to_string(index= False)}\n")
            else:
                print(f"There are no signups waiting.")
            confirm = input("Confirm (y/n):")

    if chosen.shape[0]:
        chosen[NINF[-2]] = [UId +i for i in range(1, chosen.shape[0]+1)]
        UId = chosen[NINF[-2]].max()
        Chosen = Chosen.append(chosen, ignore_index=True)

# The days and requestes are saved
Chosen.to_csv(chosenF, index= False)
unflagIt(flag)

for d in DFs:
    DFs[d].to_csv(fr"{LOC}\{d}.csv", index = False)

# The SigFile is updated
if sigs:
    nxtSig = int(sigs[0].split(".")[0][3:])+1
    remove(fr"{LOC2}\{sigs[0]}")

    with open(fr"{LOC2}\Sig{nxtSig}.txt", "w") as s:
        s.write(str(UId))


import networkx as nx
from networkx.algorithms import community
from nltk.corpus import stopwords
import re

def build_graph(text):
    word_list = []
    G = nx.Graph()
    for line in text:
        line = (line.strip()).split()
        for i, word in enumerate(line):
            if i != len(line)-1:
                word_a = word
                word_b = line[i+1]
                if word_a not in word_list:
                    word_list.append(word_a)
                if word_b not in word_list:
                    word_list.append(word_b)
                if G.has_edge(word_a,word_b):
                    G[word_a][word_b]['weight'] += 1
                else:
                    G.add_edge(word_a,word_b, weight = 1)
    return G

def calculate_central_nodes(text_network):
    bc = (nx.betweenness_centrality(text_network,weight='weight'))
    nx.set_node_attributes(text_network, bc, 'betweenness')
    bc_threshold = sorted(bc.values(), reverse=True)[20]
    to_keep = [n for n in bc if bc[n] > bc_threshold]
    filtered_network = text_network.subgraph(to_keep)
    return filtered_network

def create_and_assign_communities(text_network):
    communities_generator = community.girvan_newman(text_network)
    top_level_communities = next(communities_generator)
    next_level_communities = next(communities_generator)
    return next_level_communities

def find_topics(text):
    try:
        text_network = build_graph(text)
        text_network = calculate_central_nodes(text_network)
        topics = create_and_assign_communities(text_network)
        return topics
    except:
        print("Error: Something went wrong. Check your input. You need at least 20 unique words in your text to start the analysis.")

def clean(text):
    new_text = []
    no_punct = [re.sub(r'[^\w\s]','',x) for x in text]
    stop_words = set(stopwords.words('english'))
    for line in no_punct:
        new_line = ([item.lower() for item in line.split() if not item.lower() in stop_words])
        new_text.append(' '.join((new_line)))
    return new_text


import numpy as np
from numpy import ma
import xarray as xr
from netCDF4 import Dataset
import struct
import sys
import os
import datetime as dt
import glob

"""
This module contains functions for reading external data
to use with LPT.

The data_read_function is called at various points in other LPT functions.

To add a new data set, do the following:
1) Write a read function similar to read_generic_netcdf below.
2) Add an "elif" option that calls that function in readdata
"""

################################################################################

def readdata(datetime_to_read, dataset_options_dict, verbose=None):
    """
    Main data read function. Get data at datetime datetime_to_read.
    Based on the oprions in dataset_options_dict, it will look in the data directory
    and use the rain function specified below.

    To add a dataset type, add an elif block to this function.

    The function is expected to return a dictionary with keys 'lon', 'lat', and 'data'

    Verbose option (new 05/2023):
    - If set to None (default), it will use the verbose option from dataset_options_dict.
    - Otherwise, the value will be used *instead of* dataset_options_dict.
      This allows a function call to override the setting in dataset_options_dict.
    """

    ## Manage verbose
    if verbose is None:
        verbose_actual = dataset_options_dict['verbose']
    else:
        verbose_actual = verbose

    if dataset_options_dict['raw_data_format'] == 'generic_netcdf':
        variable_names = (dataset_options_dict['longitude_variable_name']
                , dataset_options_dict['latitude_variable_name']
                , dataset_options_dict['field_variable_name'])

        DATA = read_generic_netcdf_at_datetime(datetime_to_read
                , variable_names = variable_names
                , data_dir = dataset_options_dict['raw_data_parent_dir']
                , fmt = dataset_options_dict['file_name_format']
                , verbose = verbose_actual)

    elif dataset_options_dict['raw_data_format'] == 'generic_netcdf_with_multiple_times':
        variable_names = (dataset_options_dict['longitude_variable_name']
                , dataset_options_dict['latitude_variable_name']
                , dataset_options_dict['time_variable_name']
                , dataset_options_dict['field_variable_name'])

        DATA = read_generic_netcdf_at_datetime(datetime_to_read
                , variable_names = variable_names
                , dt_to_use = datetime_to_read
                , data_dir = dataset_options_dict['raw_data_parent_dir']
                , fmt = dataset_options_dict['file_name_format']
                , verbose = verbose_actual)

    elif dataset_options_dict['raw_data_format'] == 'cmorph':
        DATA = read_cmorph_at_datetime(datetime_to_read
                , area = dataset_options_dict['area']
                , data_dir = dataset_options_dict['raw_data_parent_dir']
                , fmt = dataset_options_dict['file_name_format']
                , verbose = verbose_actual)

    elif dataset_options_dict['raw_data_format'] == 'imerg_hdf5':
        DATA = read_imerg_hdf5_at_datetime(datetime_to_read
                , area = dataset_options_dict['area']
                , data_dir = dataset_options_dict['raw_data_parent_dir']
                , fmt = dataset_options_dict['file_name_format']
                , verbose = verbose_actual)

    elif dataset_options_dict['raw_data_format'] == 'cfs_forecast':
        fcst_hour = int((datetime_to_read - dataset_options_dict['datetime_init']).total_seconds()/3600)
        fcst_resolution_hours = dataset_options_dict['data_time_interval']
        if fcst_hour < 1: # There is no data in the file for fcst = 0. Use 6h fcst values.
            records = [1,]
        else:
            records = [int(fcst_hour/fcst_resolution_hours),]

        DATA = read_cfs_rt_at_datetime(dataset_options_dict['datetime_init'] # datetime_to_read
                , data_dir = dataset_options_dict['raw_data_parent_dir']
                , fmt = dataset_options_dict['file_name_format']
                , records = records
                , verbose = verbose_actual)
        DATA['data'] = ma.masked_array(DATA['precip'][0])

    ## -- Add an elif block here for new datasets. --

    else:
        print(('ERROR! '+dataset_options_dict['raw_data_format'] + ' is not a valid raw_data_format!'), flush=True)
        DATA = None

    return DATA

################################################################################
## Read functions for generic NetCDF data.
################################################################################

def read_generic_netcdf(fn, variable_names=('lon','lat','rain'), dt_to_use=None):
    """
    DATA = read_generic_netcdf(fn)

    output is like this:
    list(DATA)
    Out[12]: ['lon', 'lat', 'precip']
    In [21]: DATA['lon'].shape
    Out[21]: (1440,)
    In [22]: DATA['lat'].shape
    Out[22]: (400,)
    In [23]: DATA['precip'].shape
    Out[23]: (400, 1440)
    """

    DATA = {}
    with xr.open_dataset(fn) as DS:
        DATA['lon'] = DS[variable_names[0]].values
        DATA['lat'] = DS[variable_names[1]].values
        ## If no time variable, just retrieve the 2-D data as it is.
        if not dt_to_use is None: #'time' in list(DS.variables):
            DATA['data'] = DS.sel({variable_names[2]:str(dt_to_use)},method='nearest')[variable_names[3]].values
        else:
            DATA['data'] = DS[variable_names[2]].values

    DATA['data'] = np.ma.masked_array(DATA['data'], mask=np.isnan(DATA['data']))

    ## Need to get from (-180, 180) to (0, 360) longitude.
    lon_lt_0, = np.where(DATA['lon'] < -0.0001)
    lon_ge_0, = np.where(DATA['lon'] > -0.0001)
    if len(lon_lt_0) > 0:
        DATA['lon'][lon_lt_0] += 360.0
        DATA['lon'] = np.concatenate((DATA['lon'][lon_ge_0], DATA['lon'][lon_lt_0]))
        DATA['data'] = np.concatenate((DATA['data'][:,lon_ge_0], DATA['data'][:,lon_lt_0]), axis=1)

    return DATA

def read_generic_netcdf_at_datetime(dt, data_dir='.'
        , variable_names=('lon','lat','rain'), dt_to_use=None, fmt='gridded_rain_rates_%Y%m%d%H.nc'
        , verbose=False):

    fn = (data_dir + '/' + dt.strftime(fmt))
    DATA=None

    if not os.path.exists(fn):
        print('File not found: ', fn)
    else:
        if verbose:
            print(fn)
        DATA=read_generic_netcdf(fn,
            variable_names = variable_names,
            dt_to_use = dt_to_use)

    return DATA

################################################################################
## Read functions for specific datasets.
################################################################################

"""
CMORPH reading functions.
"""
def read_cmorph_rt_bin(fn, area=[0,360,-90,90]):

    """
    DATA = read_cmorph_rt_bin(fn)
    DATA is a dict with keys lon, lat, and precip.

    CMORPH RT files are binary.
    The GrADS control file below is used as the basis for this function:

    DSET ^../%y4/%y4%m2/CMORPH_V0.x_RT_8km-30min_%y4%m2%d2%h2
    OPTIONS little_endian template
    UNDEF -999.0
    TITLE CMORPH Rain Rate (Real-Time Version)
    XDEF  4948 LINEAR   0.0363783345 0.072756669
    YDEF  1649 LINEAR -59.963614312  0.072771376
    ZDEF     1 LEVELS   1
    TDEF 99999 LINEAR 00:00z01Jan2017 30mn
    VARS 1
    cmorph  1  99  CMORPH Rain Rate [mm/hr]
    ENDVARS
    """

    dtype=np.dtype([('field1', '<i2')])
    DATA={}
    DATA['lon'] = np.arange(0.0363783345, 360.0, 0.072756669)
    DATA['lat'] = np.arange(-59.963614312, 60.0, 0.072771376)
    fid = open(fn,'rb')

    ## GrADS uses FORTRAN REAL values, which is np.float32 for Python.
    DATA['data'] = np.fromfile(fid, dtype=np.float32, count=2*4948*1649)
    if sys.byteorder == 'big': # Data is little endian.
        DATA['data'] = DATA['data'].byteswap()

    ## Shape and scale the data.
    DATA['data'] = np.reshape(np.double(DATA['data']), [2, 1649, 4948])
    DATA['data'][DATA['data'] < -0.001] = 0.0 # Usually, missing high latitude data.
    fid.close()

    ## Cut out area.
    keep_lon, = np.where(np.logical_and(DATA['lon'] > area[0], DATA['lon'] < area[1]))
    keep_lat, = np.where(np.logical_and(DATA['lat'] > area[2], DATA['lat'] < area[3]))

    DATA['lon'] = DATA['lon'][keep_lon[0]:keep_lon[-1]+1]
    DATA['lat'] = DATA['lat'][keep_lat[0]:keep_lat[-1]+1]
    DATA['data'] = DATA['data'][:, keep_lat[0]:keep_lat[-1]+1, keep_lon[0]:keep_lon[-1]+1]
    DATA['data'] = 0.5*(DATA['data'][0,:,:] + DATA['data'][1,:,:])

    return DATA

def read_cmorph_at_datetime(dt_this, force_rt=False, data_dir='.'
        , fmt='CMORPH_V0.x_RT_8km-30min_%Y%m%d%H'
        , verbose=False, area=[0,360,-90,90]):

    """
    DATA = read_cmorph_at_datetime(dt, force_rt=False, verbose=False)

    DATA is a dict with keys lon, lat, and precip.

    Based on the provided datetime dt, read in the CMORPH data.

    By default, it will first check for the research product,
    and use the realtime product if the research product was not found.
    However, if force_rt = True, it just uses the realtime product.
    """

    ## First try research product
    fn = (data_dir + '/' + dt_this.strftime(fmt))
    if verbose:
        print(fn)
    DATA = read_cmorph_rt_bin(fn, area=area)

    DATA['data'] = ma.masked_array(DATA['data'])
    return DATA

def read_imerg_hdf5_at_datetime(dt_this, force_rt=False, data_dir='.'
        , fmt='%Y/%m/%d/3B-HHR.MS.MRG.3IMERG.%Y%m%d-S%H*.HDF5'
        , verbose=False, area=[0,360,-90,90]):

    """
    DATA = read_imerg_hdf5_at_datetime(dt_this, force_rt=False, data_dir='.'
        , fmt='%Y/%m/%d/3B-HHR.MS.MRG.3IMERG.%Y%m%d-S%H*.HDF5'
        , verbose=False, area=[0,360,-90,90])

    DATA is a dict with keys lon, lat, and precip.

    Based on the provided datetime dt, read in the IMERG HDF data.

    By default, it will first check for the final product,
    and use the "late" realtime product if the final product was not found.
    However, if force_rt = True, it just uses the "late" realtime product.

    (It will search for a filename with modified fmt to check for "late" product
    - append 'late/' to the front of the directory path.
    - replace '3B-HHR' with '3B-HHR-L').
    """

    fn_list = sorted(glob.glob(data_dir + '/' + dt_this.strftime(fmt)))
    if len(fn_list) < 1:
        if not force_rt:
            ## Try "late" realtime data.
            print('Final data version not found. Trying to use late realtime data instead.')
            fmt_rt = 'late/' + fmt.replace('3B-HHR','3B-HHR-L')
            fn_list = sorted(glob.glob(data_dir + '/' + dt_this.strftime(fmt_rt)))

    if len(fn_list) < 1:
        print('WARNING: No input data found.')

    fn = fn_list[0]
    if verbose:
        print(fn)

    with Dataset(fn) as DS:
        lon_rain = DS['Grid']['lon'][:]
        lat_rain = DS['Grid']['lat'][:]
        rain = DS['Grid']['precipitationCal'][:][0].T

    if len(fn_list) > 1:
        fn = fn_list[1]
        if verbose:
            print(fn)

        with Dataset(fn) as DS:
            rain30 = DS['Grid']['precipitationCal'][:][0].T

        rain = 0.5 * (rain + rain30)

    ## lon -180:0 --> 180:360
    idx_neg_lon = [x for x in range(len(lon_rain)) if lon_rain[x] < -0.0001]
    idx_pos_lon = [x for x in range(len(lon_rain)) if lon_rain[x] > -0.0001]

    lon_rain = np.append(lon_rain[idx_pos_lon[0]:idx_pos_lon[-1]+1], 360.0 + lon_rain[idx_neg_lon[0]:idx_neg_lon[-1]+1], axis=0)
    rain = np.append(rain[:,idx_pos_lon[0]:idx_pos_lon[-1]+1], rain[:,idx_neg_lon[0]:idx_neg_lon[-1]+1], axis=1)

    DATA={}
    DATA['lon'] = lon_rain
    DATA['lat'] = lat_rain
    DATA['data'] = ma.masked_array(rain)

    ## Cut out area.
    keep_lon, = np.where(np.logical_and(DATA['lon'] > area[0], DATA['lon'] < area[1]))
    keep_lat, = np.where(np.logical_and(DATA['lat'] > area[2], DATA['lat'] < area[3]))

    DATA['lon'] = DATA['lon'][keep_lon[0]:keep_lon[-1]+1]
    DATA['lat'] = DATA['lat'][keep_lat[0]:keep_lat[-1]+1]
    DATA['data'] = DATA['data'][keep_lat[0]:keep_lat[-1]+1, keep_lon[0]:keep_lon[-1]+1]

    return DATA

################################################################################
################################################################################
################################################################################

"""
CFS Grib2 reading function
"""

def read_cfs_rt_at_datetime(dt_this, data_dir = './'
                , fmt = 'cfs.%Y%m%d/%H/time_grib_01/prate.01.%Y%m%d%H.daily.grb2'
                , records=range(1,45*4+1), verbose=False):

    fn = (data_dir + '/' + dt_this.strftime(fmt))
    if verbose:
        print(fn, flush=True)

    return read_cfs_rt_grib2(fn, records=records, verbose=verbose)

def read_cfs_rt_grib2(fn, records=range(1,45*4+1), verbose=False):
    """
    RT = read_cfs_rt_grib2(fn, records=N)

    N is the list of records to get.
    By default, get the first 45 days, 6 hourly intervals.

    example output:
    In [23]: RT['lon'].shape
    Out[23]: (384,)

    In [24]: RT['lat'].shape
    Out[24]: (190,)

    In [25]: RT['precip'].shape
    Out[25]: (180, 190, 384)
    """

    import gdal  # Import gdal if dealing with grib data.

    DS = gdal.Open(fn, gdal.GA_ReadOnly)
    width = DS.RasterXSize
    height = DS.RasterYSize
    lon =  np.arange(0.0,359.062 + 0.5,0.938)
    ## grid file with Gaussian latitude was obtained from wgrib2 like this:
    ## wgrib2 -d 1 -gridout grid.txt /home/orca/data/model_fcst_grib/cfs/cfs.20190508/00/time_grib_01/prate.01.2019050800.daily.grb2
    ## awk -F, '{print $3}' grid.txt | uniq | tr "\n" ", "
    lat = np.flip(np.array([-89.277, -88.340, -87.397, -86.454, -85.509
                , -84.565, -83.620, -82.676, -81.731, -80.786
                , -79.841, -78.897, -77.952, -77.007, -76.062
                , -75.117, -74.173, -73.228, -72.283, -71.338
                , -70.393, -69.448, -68.503, -67.559, -66.614
                , -65.669, -64.724, -63.779, -62.834, -61.889
                , -60.945, -60.000, -59.055, -58.110, -57.165
                , -56.220, -55.275, -54.330, -53.386, -52.441
                , -51.496, -50.551, -49.606, -48.661, -47.716
                , -46.771, -45.827, -44.882, -43.937, -42.992
                , -42.047, -41.102, -40.157, -39.212, -38.268
                , -37.323, -36.378, -35.433, -34.488, -33.543
                , -32.598, -31.653, -30.709, -29.764, -28.819
                , -27.874, -26.929, -25.984, -25.039, -24.094
                , -23.150, -22.205, -21.260, -20.315, -19.370
                , -18.425, -17.480, -16.535, -15.590, -14.646
                , -13.701, -12.756, -11.811, -10.866, -9.921
                , -8.976, -8.031, -7.087, -6.142, -5.197
                , -4.252, -3.307, -2.362, -1.417, -0.472
                , 0.472, 1.417, 2.362, 3.307, 4.252
                , 5.197, 6.142, 7.087, 8.031, 8.976
                , 9.921, 10.866, 11.811, 12.756, 13.701
                , 14.646, 15.590, 16.535, 17.480, 18.425
                , 19.370, 20.315, 21.260, 22.205, 23.150
                , 24.094, 25.039, 25.984, 26.929, 27.874
                , 28.819, 29.764, 30.709, 31.653, 32.598
                , 33.543, 34.488, 35.433, 36.378, 37.323
                , 38.268, 39.212, 40.157, 41.102, 42.047
                , 42.992, 43.937, 44.882, 45.827, 46.771
                , 47.716, 48.661, 49.606, 50.551, 51.496
                , 52.441, 53.386, 54.330, 55.275, 56.220
                , 57.165, 58.110, 59.055, 60.000, 60.945
                , 61.889, 62.834, 63.779, 64.724, 65.669
                , 66.614, 67.559, 68.503, 69.448, 70.393
                , 71.338, 72.283, 73.228, 74.173, 75.117
                , 76.062, 77.007, 77.952, 78.897, 79.841
                , 80.786, 81.731, 82.676, 83.620, 84.565
                , 85.509, 86.454, 87.397, 88.340, 89.277]), axis=0)

    num_list = []
    for band in records:
        if verbose:
            print('Record #' + str(band), flush=True)
        data_array = DS.GetRasterBand(band).ReadAsArray()
        for row in data_array:
            for value in row:
                num_list.append(value*3600.0) # kg/m2/sec --> mm/h

    DS = None # Close the file.

    precip = np.array(num_list).reshape([len(records), len(lat), len(lon)])

    DATA={}
    DATA['lon'] = lon
    DATA['lat'] = lat
    DATA['precip'] = precip

    return DATA

def read_cfsr_grib2(fn, band_list=None, verbose=False):
    """
    RT = read_cfsr_grib2(fn)

    example output:
    In [23]: RT['lon'].shape
    Out[23]: (384,)

    In [24]: RT['lat'].shape
    Out[24]: (190,)

    In [25]: RT['precip'].shape
    Out[25]: (180, 190, 384)
    """

    DS = gdal.Open(fn, gdal.GA_ReadOnly)
    width = DS.RasterXSize
    height = DS.RasterYSize
    lon =  np.arange(0.0,359.51,0.5)
    lat =  np.arange(90.0,-90.01,-0.5)
    n_records = DS.RasterCount

    num_list = []

    if band_list is None:
        band_list = range(1, n_records+1)

    for band in band_list:
        if verbose:
            print((str(band) + ' of ' + str(n_records)))
        data_array = DS.GetRasterBand(band).ReadAsArray()
        for row in data_array:
            for value in row:
                num_list.append(value)

    DS = None # Close the file.

    precip = np.array(num_list).reshape([int(len(band_list)/6), 6, len(lat), len(lon)])
    #precip /= 1e6  # Values in file are multiplied by 1e6.
                   # kg/m2 in 1h is equivalent to mm/h.

    DATA={}
    DATA['lon'] = lon
    DATA['lat'] = lat
    DATA['precip'] = precip

    return DATA

def get_cfsr_6h_rain(dt_ending, verbose=False):

    """
    Read in the rainfall using read_cfs_historical_grib2(fn)
    Then calculate the 6 hourly rain rate (mm/h) and return it.

    CFSR rain is stored in monthly files. It it initialized every 6 h,
    and the data provide hourly accumulations (in kg/m^2, equivalent to mm) like this:

    1:0:d=2011120100:APCP:surface:0-1 hour acc fcst:
    2:94325:d=2011120100:APCP:surface:1-2 hour acc fcst:
    3:193206:d=2011120100:APCP:surface:2-3 hour acc fcst:
    4:309596:d=2011120100:APCP:surface:3-4 hour acc fcst:
    5:421187:d=2011120100:APCP:surface:4-5 hour acc fcst:
    6:537704:d=2011120100:APCP:surface:5-6 hour acc fcst:

    To get the 6 hourly accumulation, all 6 of these need to be added.
    Then take the mean (e.g., divide by 6h) to get mm/h.
    """

    dt_beginning = dt_ending - dt.timedelta(hours=6)

    if dt_beginning < dt.datetime(2011,3,31,23,59,0):
        fn_beginning = ('/home/orca/data/model_anal/cfsr/rain_accum/' + dt_beginning.strftime('%Y')
            + '/apcp.gdas.' + dt_beginning.strftime('%Y%m') + '.grb2')
    else:
        fn_beginning = ('/home/orca/data/model_anal/cfsr/rain_accum/' + dt_beginning.strftime('%Y')
            + '/apcp.cdas1.' + dt_beginning.strftime('%Y%m') + '.grb2')

    if verbose:
        print(fn_beginning, flush=True)

    rec_num = 1 + int((dt_beginning - dt.datetime(dt_beginning.year, dt_beginning.month,1,0,0,0)).total_seconds()/3600.0)
    F = read_cfsr_grib2(fn_beginning, band_list=range(rec_num,rec_num+6,1), verbose=verbose)

    precip6hr = np.nanmean(F['precip'], axis=1)[0]

    DATA={}
    DATA['lon'] = F['lon']
    DATA['lat'] = F['lat']
    DATA['precip'] = precip6hr

    return DATA