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

import platform
from time import sleep
import logging
import queue
import fnmatch
import os
import glob
from .storagemanager import StorageManager, StorageManagerSQLite
from .mpreaderwriter import DockingFileReader
from .mpreaderwriter import Writer
from .exceptions import MultiprocessingError
import traceback
from datetime import datetime

os_string = platform.system()
if os_string == "Darwin":  # mac
    import multiprocess as multiprocessing
else:
    import multiprocessing

class MPManager:
    def __init__(
        self,
        storageman,
        storageman_class=StorageManagerSQLite,
        mode="dlg",
        chunk_size=1,
        max_poses=3,
        interaction_tolerance=None,
        store_all_poses=False,
        target=None,
        add_interactions=False,
        interaction_cutoffs=[3.7, 4.0],
        receptor_file=None,
        file_sources={
            "file": [[]],
            "file_path": {"path": [[]], "pattern": "*.dlg*", "recursive": None},
            "file_list": [[]],
        },
        file_pattern="*.dlg*",
        max_proc=None,
    ):

        # confirm that requested parser mode is implemented
        self.implemented_modes = ["dlg", "vina"]
        if mode not in self.implemented_modes:
            raise NotImplementedError(
                f"Requested file parsing mode {mode} not yet implemented"
            )
        self.mode = mode
        self.chunk_size = chunk_size
        self.max_poses = max_poses
        self.store_all_poses = store_all_poses
        self.interaction_tolerance = interaction_tolerance
        self.target = target
        self.add_interactions = add_interactions
        self.interaction_cutoffs = interaction_cutoffs
        self.receptor_file = receptor_file
        self.file_sources = file_sources
        self.file_pattern = file_pattern

        self.storageman = storageman
        self.storageman_class = storageman_class
        self.num_files = 0
        self.max_proc = max_proc

    def process_files(self):
        if self.max_proc is None:
            self.max_proc = multiprocessing.cpu_count()
        self.num_readers = self.max_proc - 1
        self.queueIn = multiprocessing.Queue(maxsize=2 * self.max_proc)
        self.queueOut = multiprocessing.Queue(maxsize=2 * self.max_proc)
        # start the workers in background
        self.workers = []
        self.p_conn, self.c_conn = multiprocessing.Pipe(True)
        logging.info("Starting {0} file readers".format(self.num_readers))
        for i in range(self.num_readers):
            # one worker is started for each processor to be used
            s = DockingFileReader(
                self.queueIn,
                self.queueOut,
                self.c_conn,
                self.storageman,
                self.storageman_class,
                self.mode,
                self.max_poses,
                self.interaction_tolerance,
                self.store_all_poses,
                self.target,
                self.add_interactions,
                self.interaction_cutoffs,
                self.receptor_file,
            )
            # this method calls .run() internally
            s.start()
            self.workers.append(s)

        # start the writer to process the data from the workers
        w = Writer(
            self.queueOut,
            self.num_readers,
            self.c_conn,
            self.chunk_size,
            self.storageman,
            self.mode,
        )

        w.start()
        self.workers.append(w)

        # process items in the queue
        try:
            self._process_sources()
        except Exception as e:
            tb = traceback.format_exc()
            self._kill_all_workers(e, "file sources", tb)
        # put as many poison pills in the queue as there are workers
        for i in range(self.num_readers):
            self.queueIn.put(None)

        # check for exceptions
        while w.is_alive():
            sleep(0.5)
            self._check_for_worker_exceptions()

        w.join()

        logging.info("Wrote {0} files to database".format(self.num_files))

    def _process_sources(self):
        # add individual file(s)
        if self.file_sources["file"] != [[]]:
            for file_list in self.file_sources["file"]:
                for file in file_list:
                    if (
                        fnmatch.fnmatch(file, self.file_pattern)
                        and file != self.receptor_file
                    ):
                        self._add_to_queue(file)
        # add files from file path(s)
        if self.file_sources["file_path"]["path"] != [[]]:
            for path_list in self.file_sources["file_path"]["path"]:
                for path in path_list:
                    # scan for ligand dlgs
                    for files in self._scan_dir(
                        path, self.file_pattern, recursive=True
                    ):
                        for f in files:
                            self._add_to_queue(f)
        # add files from file list(s)
        if self.file_sources["file_list"] != [[]]:
            for filelist_list in self.file_sources["file_list"]:
                for filelist in filelist_list:
                    self._scan_file_list(filelist, self.file_pattern.replace("*", ""))

    def _add_to_queue(self, file):
        max_attempts = 750
        timeout = 0.5  # seconds
        if self.receptor_file is not None:
            if (
                os.path.split(file)[-1] == os.path.split(self.receptor_file)[-1]
            ):  # check that we don't try to add the receptor
                return
        attempts = 0
        while True:
            if attempts >= max_attempts:
                raise MultiprocessingError(
                    "Something is blocking the progressing of file reading. Exiting program."
                ) from queue.Full
            try:
                self.queueIn.put(file, block=True, timeout=timeout)
                self.num_files += 1
                self._check_for_worker_exceptions()
                break
            except queue.Full:
                # logging.debug(f"Queue full: queueIn.put attempt {attempts} timed out. {max_attempts - attempts} put attempts remaining.")
                attempts += 1
                self._check_for_worker_exceptions()

    def _check_for_worker_exceptions(self):
        if self.p_conn.poll():
            error, tb, filename = self.p_conn.recv()
            logging.error(f"Caught error in multiprocessing from {filename}")
            # don't kill parser errors, only database error
            if filename == "Database":
                self._kill_all_workers(error, filename, tb)
            else:
                with open("ringtail_failed_files.log", 'a') as f:
                    f.write(str(datetime.now()) + f"\tRingtail failed to parse {filename}\n")
                    logging.debug(tb)

    def _kill_all_workers(self, error, filename, tb):
        for s in self.workers:
            s.kill()
        logging.debug(f"Error encountered while handling {filename}")
        logging.debug(tb)
        raise error

    def _scan_dir(self, path, pattern, recursive=False):
        """scan for valid output files in a directory
        the pattern is used to glob files
        optionally, a recursive search is performed
        """
        logging.info(
            "-Scanning directory [%s] for files (pattern:|%s|)" % (path, pattern)
        )
        if recursive:
            path = os.path.normpath(path)
            path = os.path.expanduser(path)
            for dirpath, dirnames, filenames in os.walk(path):
                yield (  # <----
                    os.path.join(dirpath, f)
                    for f in fnmatch.filter(filenames, "*" + pattern)
                )
        else:
            yield glob(os.path.join(path, pattern))  # <----

    def _scan_file_list(self, filename, pattern=".dlg"):
        """read file names from file list"""
        lig_accepted = []
        c = 0
        with open(filename, "r") as fp:
            for line in fp.readlines():
                line = line.strip()
                c += 1
                if os.path.isfile(line):
                    if line.endswith(pattern) or line.endswith(pattern + ".gz"):
                        lig_accepted.append(line)
                else:
                    logging.warning("Warning! file |%s| does not exist" % line)
        if len(lig_accepted) == 0:
            raise MultiprocessingError(
                "*ERROR* No valid files were found when reading from |%s|" % filename
            )
        logging.info(
            "# [ %5.3f%% files in list accepted (%d) ]"
            % ((len(lig_accepted) / c * 100, c))
        )

        for file in lig_accepted:
            if file != self.receptor_file:
                self._add_to_queue(file)


from turtle import Turtle, Screen, colormode
from runner import Runner
from cars import Cars
from scoreboard import Scoreboard
import time
from random import randint

colormode(255)

screen = Screen()
screen.setup(width=600, height=600)
screen.bgcolor('white')
screen.title('Turtle Passing Capstone')
screen.tracer(0)

pet = Runner()
cars = Cars()
scoreboard = Scoreboard()

screen.listen()
screen.onkey(pet.move, 'Up')
screen.onkey(pet.move, 'w')
screen.onkey(pet.move, '8')
screen.onkey(pet.move, 'space')

game_is_on = True
while game_is_on:
    time.sleep(cars.car_speed)
    screen.update()
    cars.create_car()
    cars.move_cars()

    # Detect collision with cars
    for car in cars.all_cars:
        if car.distance(pet) < 20:
            game_is_on = False
            scoreboard.game_over()

    if pet.ycor() > 280:
        pet.come_back()
        cars.speed_up()
        scoreboard.level_up()

screen.exitonclick()


class Solution:
    def isPalindrome(self, s: str) -> bool:
        def helper(sub):
            return(ord(sub) >= ord('a') and ord(sub) <= ord('z')) or (ord(sub) >= ord('0') and ord(sub) <= ord('9'))
        l,r = 0,len(s)-1
        while l < r:
            while l<r and not helper(s[l].lower()):
                l +=1
            while l < r and not helper(s[r].lower()):
                r -= 1
            if s[l].lower() == s[r].lower():
                l += 1
                r -= 1
            else:
                return False

        return True


import numpy as np
import matplotlib.pyplot as plt

x1 = np.array([0.5, 0.1, 0.2, 0.4, 0.3, 0.2, 0.2, 0.1, 0.35, 0.25])[:,None]
x2 = np.array([0.9, 0.8, 0.75, 1.0])[:,None]
N = len(x1) + len(x2)
prior1 = len(x1) / N
prior2 = len(x2) / N

def gaussian(mean, var, xx):
    return np.exp(-0.5 * (xx - mean)**2 / var) * 1/np.sqrt(2*np.pi*var)

def plot_tutorial_3a():
    plt.clf()

    """Code for plotting the fitted Gaussians"""
    # X_grid = np.arange(-.1, 1.2, 0.01)[:,None]
    # gaussian_1 = gaussian(x1.mean(), x1.var(), X_grid)
    # gaussian_2 = gaussian(x2.mean(), x2.var(), X_grid)
    # plt.plot(X_grid, gaussian_1, label='Gaussian 1')
    # plt.plot(X_grid, gaussian_2, label='Gaussian 2')

    """ Code to plot the scores p(x,y) = P(y)p(x|y) of the given datapoints """
    xx = np.concatenate((x1,x2))
    p_x_given_class_1 = prior1 * gaussian(x1.mean(), x1.var(), xx)
    plt.scatter(xx, p_x_given_class_1, label='score for data given class 1')
    p_x_given_class_2 = prior2 * gaussian(x2.mean(), x2.var(), xx)
    plt.scatter(xx, p_x_given_class_2, label='score for data given class 2')

    """ Code to plot the scores p(x,y) = P(y)p(x|y) as a function of x"""
    xx = np.linspace(0, 1, num=100)
    p_x_given_class_1 = prior1 * gaussian(x1.mean(), x1.var(), xx)
    plt.plot(xx, p_x_given_class_1, label='score for class 1')
    p_x_given_class_2 = prior2 * gaussian(x2.mean(), x2.var(), xx)
    plt.plot(xx, p_x_given_class_2, label='score for class 2')

    # change labels and add legend
    plt.xlabel('location x')
    plt.ylabel('p(x,y)')
    plt.legend()

    # Show the plot
    plt.show()

def plot_gaussians_only():
    """Code for plotting only the Gaussians"""
    X_grid = np.arange(-.1, 1.2, 0.01)[:,None]
    gaussian_1 = gaussian(x1.mean(), x1.var(), X_grid)
    gaussian_2 = gaussian(x2.mean(), x2.var(), X_grid)

    plt.plot(X_grid, gaussian_1, label='Gaussian 1')
    plt.scatter(x1, np.zeros_like(x1), label='points x1')
    plt.plot(X_grid, gaussian_2, label='Gaussian 2')
    plt.scatter(x2, np.zeros_like(x2), label='points x2')
    plt.legend()
    plt.show()

def plot_probabilities_only(x):
    """Code for plotting only the scores p(x|y)"""
    p_x_given_class_1 = prior1 * gaussian(x1.mean(), x1.var(), x)
    plt.scatter(x, p_x_given_class_1, label='class 1')
    p_x_given_class_2 = prior2 * gaussian(x2.mean(), x2.var(), x)
    plt.scatter(x, p_x_given_class_2, label='class 2')
    plt.legend()
    plt.show()

def probability_of_class_1_given_x(x: float) -> float:
    # P(x=0.6) = P(x=0.6 | y=1)P(y=1) + P(x=0.6 | y=2)P(y=2)
    probability_of_x = gaussian(x1.mean(), x1.var(), x) * prior1 + gaussian(x2.mean(), x2.var(), x) * prior2

    # P(y=1 | x=0.6) = P(x=0.6 | y=1)P(y=1) / P(x=0.6)
    probability_asked = gaussian(x1.mean(), x1.var(), x) * prior1 / probability_of_x
    print(f'P(y=1|x={x}) = {probability_asked}')

probability_of_class_1_given_x(0.6)
plot_tutorial_3a()


import csv
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn import metrics
import pickle

def write_csv_file(path, data_array,access_type):
    with open(path, access_type, newline='', encoding="utf-8") as f:
        writer = csv.writer(f)
        for row in data_array:
            writer.writerow(row)

def RandomForest():
    tfidf_value_matrix = pd.read_csv('./csv/5_tfidf_for_corpus.csv')

    number_of_columns = len(tfidf_value_matrix.columns)

    #tfidf value matrix(independent variables)
    X = tfidf_value_matrix.iloc[:, 0:number_of_columns-2].values
    #get story point list (dependent variable)
    y = tfidf_value_matrix.iloc[:, number_of_columns-1].values

    #separate training set and test set
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)

    sc = StandardScaler()
    X_train = sc.fit_transform(X_train)
    X_test = sc.transform(X_test)

    regressorModel = RandomForestRegressor(n_estimators=20, random_state=0)
    regressorModel.fit(X_train, y_train)

    ##save the model
    pickle.dump(regressorModel, open('./Models/RandomForestModel.sav', 'wb'))

    #Evaluation
    y_pred = regressorModel.predict(X_test)

    print('Mean Absolute Error:', metrics.mean_absolute_error(y_test, y_pred))
    print('Mean Squared Error:', metrics.mean_squared_error(y_test, y_pred))
    print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test, y_pred)))

    accuracy={
        'Mean Absolute Error': metrics.mean_absolute_error(y_test, y_pred),
        'Mean Squared Error':metrics.mean_squared_error(y_test, y_pred),
        'Root Mean Squared Error': np.sqrt(metrics.mean_squared_error(y_test, y_pred))
    }

    ###estimates story point for the dataset (used to train)

    ##read and estimate for new dataset
    dataset = pd.read_csv('./csv/5_tfidf_for_corpus.csv')

    Existing_text_features=sc.fit_transform(dataset.iloc[:, 0:number_of_columns-2].values)

    text_score = regressorModel.predict(Existing_text_features)

    #add new colums to write new values
    dataset["test_score"]=text_score
    dataset["test_score_round"]=np.rint(text_score)

    write_csv_file('./csv/6_estimated_new_text_score.csv',dataset.to_numpy(),'w')

    ##save resultant format
    resultant_text_score = dataset[['issue_key','story_point','test_score', 'test_score_round']]
    write_csv_file('./csv/7_resultant_new_text_score.csv',[['Issue key','previous_story_point','estimated_test_score', 'text_score_rounded']],'w')
    write_csv_file('./csv/7_resultant_new_text_score.csv',resultant_text_score.to_numpy(),'a')

    ###refactor dataset 2 for linear regressor####
    input_dataset= pd.read_csv("./csv/2_input_dataset.csv")
    estimated_text_score_data = pd.read_csv("./csv/7_resultant_new_text_score.csv")

    ##merge 2 dataset by issue key
    merged_dataset=pd.merge(input_dataset, estimated_text_score_data, on='Issue key')
    merged_dataset.to_csv('./csv/8_merged_datset.csv', encoding='utf-8', mode='w', header=True, index=False)
    input_dataset2=merged_dataset[['Issue key','previous_story_point','estimated_test_score','Number of developers','Number of comments']]
    input_dataset2.to_csv('./csv/9_input_datset2.csv', encoding='utf-8', mode='w', header=True, index=False)

    return accuracy

def e():
  return 4;


import csv
import math
import random
import sys

class Advertiser:
    def __init__(self, id, keyword, bidValue, budget):
        self.id = id
        self.keyword = keyword
        self.bidValue = bidValue
        self.budget = budget

# Read queries and create dictionaries for advertisers and queries
def readData():
    queryNeighbors = dict()
    advertiserBudgets = dict()
    advertisersObjects = dict()
    with open('bidder_dataset.csv', 'r') as bidderDataset:
        bidderReader = csv.reader(bidderDataset, delimiter=',')
        next(bidderReader)
        for row in bidderReader:
            if row[3] != '':
                advertiserBudgets[int(row[0])] = float(row[3])
            if int(row[0]) not in advertisersObjects:
                advertisersObjects[int(row[0])] = Advertiser(int(row[0]), [row[1]], [float(row[2])],
                                                             advertiserBudgets[int(row[0])])
            else:
                advertisersObjects[int(row[0])].keyword.append(row[1])
                advertisersObjects[int(row[0])].bidValue.append(float(row[2]))
            if row[1] not in queryNeighbors:
                queryNeighbors[row[1]] = [(int(row[0]))]
            else:
                queryNeighbors[row[1]].append(int(row[0]))
    return (queryNeighbors, advertiserBudgets, advertisersObjects)

# Check if the advertiser has sufficient budget to bid for the query or not
def isNeighborAvailable(query, qNeighbors, advertisersObjects):
    for neighbor in qNeighbors:
        bidValue = advertisersObjects[neighbor].bidValue[advertisersObjects[neighbor].keyword.index(query)]
        if advertisersObjects[neighbor].budget >= bidValue:
            return True
    return False

# Greedy approach
def greedy(queries, queryNeighbors, advertiserBudgets, advertisersObjects):
    revenue = 0
    for query in queries:
        qNeighbors = queryNeighbors[query]
        if isNeighborAvailable(query, qNeighbors, advertisersObjects):
            highestBid = 0
            highestBidNeighbor = qNeighbors[0]
            for neighbor in qNeighbors:
                bidValue = advertisersObjects[neighbor].bidValue[advertisersObjects[neighbor].keyword.index(query)]
                if bidValue > highestBid and advertisersObjects[neighbor].budget >= bidValue:
                    highestBid = bidValue
                    highestBidNeighbor = neighbor
            advertisersObjects[highestBidNeighbor].budget -= highestBid
            revenue += highestBid
    return revenue

# MSVV Approach
def msvv(queries, queryNeighbors, advertiserBudgets, advertisersObjects):
    revenue = 0
    for query in queries:
        qNeighbors = queryNeighbors[query]
        if isNeighborAvailable(query, qNeighbors, advertisersObjects):
            highestValue = -sys.maxsize
            highestValueNeighbor = qNeighbors[0]
            for neighbor in qNeighbors:
                Xu = (advertiserBudgets[neighbor] - advertisersObjects[neighbor].budget) / advertiserBudgets[neighbor]
                bidValue = advertisersObjects[neighbor].bidValue[advertisersObjects[neighbor].keyword.index(query)]
                if (1 - math.exp(Xu - 1)) * bidValue > highestValue and advertisersObjects[neighbor].budget >= bidValue:
                    highestValue = (1 - math.exp(Xu - 1)) * bidValue
                    highestValueNeighbor = neighbor
                    highestBid = bidValue
            advertisersObjects[highestValueNeighbor].budget -= highestBid
            revenue += highestBid
    return revenue

# Balance Approach
def balance(queries, queryNeighbors, advertiserBudgets, advertisersObjects):
    revenue = 0
    for query in queries:
        qNeighbors = queryNeighbors[query]
        if isNeighborAvailable(query, qNeighbors, advertisersObjects):
            highestUnspentBudget = -sys.maxsize
            highestUnspentBudgetNeighbor = qNeighbors[0]
            for neighbor in qNeighbors:
                bidValue = advertisersObjects[neighbor].bidValue[advertisersObjects[neighbor].keyword.index(query)]
                if advertisersObjects[neighbor].budget > highestUnspentBudget and advertisersObjects[
                    neighbor].budget >= bidValue:
                    highestUnspentBudget = advertisersObjects[neighbor].budget
                    highestUnspentBudgetNeighbor = neighbor
                    highestBid = bidValue
            advertisersObjects[highestUnspentBudgetNeighbor].budget -= highestBid
            revenue += highestBid
    return revenue

# Main function
if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python adwords.py method_name")
        print("method_name can be either 'greedy', 'msvv' or 'balance'")
        exit()

    random.seed(0)
    revenue = 0
    totalRevenue = 0
    averageRevenue = 0
    with open('queries.txt', 'r') as q:
        queries = [query.strip() for query in q]

    queryNeighbors, advertiserBudgets, advertisersObjects = readData()
    if sys.argv[1] == 'greedy':
        totalRevenue = greedy(queries, queryNeighbors, advertiserBudgets, advertisersObjects)
        for i in range(100):
            queryNeighbors, advertiserBudgets, advertisersObjects = readData()
            tempQueries = queries
            random.shuffle(tempQueries)
            revenue += greedy(tempQueries, queryNeighbors, advertiserBudgets, advertisersObjects)
        averageRevenue = revenue / 100
    elif sys.argv[1] == 'msvv':
        totalRevenue = msvv(queries, queryNeighbors, advertiserBudgets, advertisersObjects)
        for i in range(100):
            queryNeighbors, advertiserBudgets, advertisersObjects = readData()
            tempQueries = queries
            random.shuffle(tempQueries)
            revenue += msvv(tempQueries, queryNeighbors, advertiserBudgets, advertisersObjects)
        averageRevenue = revenue / 100
    elif sys.argv[1] == 'balance':
        totalRevenue = balance(queries, queryNeighbors, advertiserBudgets, advertisersObjects)
        for i in range(100):
            queryNeighbors, advertiserBudgets, advertisersObjects = readData()
            tempQueries = queries
            random.shuffle(tempQueries)
            revenue += balance(tempQueries, queryNeighbors, advertiserBudgets, advertisersObjects)
        averageRevenue = revenue / 100
    else:
        print('Invalid method. Please try again.')
        exit()

    print('Revenue:', round(totalRevenue, 2))
    optimalMatching = sum(advertiserBudgets.values())
    competitiveRatio = averageRevenue / optimalMatching
    print('Competitive Ratio:', round(competitiveRatio, 2))

P5 - AdWords via Graph Matching/adwords.py


# -*- coding: utf-8 -*-

"""SNMP Experiments."""

import asyncio
import warnings

from pysnmp import hlapi
from pysnmp.entity import engine, config
from pysnmp.carrier.asyncio.dgram import udp
from pysnmp.entity.rfc3413 import ntfrcv
from pysnmp.smi import builder, compiler, view, rfc1902

DEFAULT_ADDRESSS = "0.0.0.0"
DEFAULT_PORT = 162
DEFAULT_COMMUNITY = "public"
DEFAULT_MIBS = ("SNMPv2-MIB", "IF-MIB", "XYTRONIX-MIB")  # Must be a tuple not a list.

_view_controller = None

def _make_object(*id_parts):
    """Construct a pySNMP ObjectType from the given identity."""
    return hlapi.ObjectType(hlapi.ObjectIdentity(*id_parts))

def _make_get(address, community, *objects, port=161, mp_model=1):
    """Construct a pySNMP get command.

    Defaults to SNMPv2c and port 161.

    """
    engine = hlapi.SnmpEngine()
    community = hlapi.CommunityData(community, mpModel=mp_model)
    target = hlapi.UdpTransportTarget((address, port))
    context = hlapi.ContextData()
    return hlapi.getCmd(engine, community, target, context, *objects)

def _run_command(command):
    """Runs the command and returns a list of all results."""
    # Implicitly calls next() repeatdly on the command iterator until there is nothing left.
    # Supports both single GETs as well as GETBULKs.
    return [result for result in command]

def _extract_errors(results):
    """Returns a list of error tuples if any errors exist in the results.

    Each error tuple consists of (error_indication, error_status, error_index).

    """
    errors = []
    for error_indication, error_status, error_index, var_binds in results:
        if error_indication or error_status:
            error_text = error_status.prettyPrint() if error_status else error_status
            object_id = var_binds[int(error_index) - 1][0] if error_index else "?"
            errors.append((error_indication, error_text, object_id))
    return errors

def _extract_values(results):
    """Returns a dict of data values from the results."""
    values = {}
    for result in results:
        var_binds = result[-1]
        for var_bind in var_binds:
            values[var_bind[0].prettyPrint()] = var_bind[1]
    return values

def _print_errors(errors):
    """Prints errors to the screen."""
    for error_indication, error_text, object_id in errors:
        if error_indication:
            print(error_indication)
        elif error_text:
            print(f"{error_text} at {object_id}")

def _print_values(values):
    """Prints object ids and their data values to the screen."""
    for object_id, value in values.items():
        print(object_id, "=", value)

def _print_results(results):
    """Prints the results of snmp commands and/or any related errors."""
    _print_errors(_extract_errors(results))
    _print_values(_extract_values(results))

def _old_print_results(command):
    """Prints the results of snmp commands and/or any related errors. (Deprecated)"""
    warnings.warn("Deprecated.", DeprecationWarning)
    errorIndication, errorStatus, errorIndex, varBinds = next(command)
    if errorIndication:
        print(errorIndication)
    elif errorStatus:
        print(
            "%s at %s"
            % (
                errorStatus.prettyPrint(),
                errorIndex and varBinds[int(errorIndex) - 1][0] or "?",
            )
        )
    else:
        for varBind in varBinds:
            print(" = ".join([x.prettyPrint() for x in varBind]))

def quickstart():
    """PySNMP Quick Start Example"""
    iterator = hlapi.getCmd(
        hlapi.SnmpEngine(),
        hlapi.CommunityData("public", mpModel=0),
        hlapi.UdpTransportTarget(("192.168.0.59", 161)),
        hlapi.ContextData(),
        hlapi.ObjectType(hlapi.ObjectIdentity("SNMPv2-MIB", "sysDescr", 0)),
    )
    errorIndication, errorStatus, errorIndex, varBinds = next(iterator)
    if errorIndication:
        print(errorIndication)
    elif errorStatus:
        print(
            "%s at %s"
            % (
                errorStatus.prettyPrint(),
                errorIndex and varBinds[int(errorIndex) - 1][0] or "?",
            )
        )
    else:
        for varBind in varBinds:
            print(" = ".join([x.prettyPrint() for x in varBind]))

def common():
    """PySNMP Tutorial Common Operations Example"""
    engine = hlapi.SnmpEngine()
    community = hlapi.CommunityData("public", mpModel=1)  # SNMPv2c
    target = hlapi.UdpTransportTarget(("192.168.0.59", 161))
    context = hlapi.ContextData()

    # object_id = hlapi.ObjectIdentity(
    #     "iso.org.dod.internet.mgmt.mib-2.system.sysDescr.0"
    # )
    sysDescr = hlapi.ObjectType(hlapi.ObjectIdentity("SNMPv2-MIB", "sysDescr", 0))
    sysUpTime = hlapi.ObjectType(hlapi.ObjectIdentity("SNMPv2-MIB", "sysUpTime", 0))
    ifInOctets = hlapi.ObjectType(hlapi.ObjectIdentity("IF-MIB", "ifInOctets", 1))
    command = hlapi.getCmd(
        engine, community, target, context, sysDescr, sysUpTime, ifInOctets
    )
    _old_print_results(command)

def temperature():
    """One-Wire Temperature sensor on ControlByWeb X-410 module."""
    temp = _make_object("XYTRONIX-MIB", "temp", 0)
    command = _make_get("192.168.0.132", "webrelay", temp)
    results = _run_command(command)
    _print_results(results)

def rewrite():
    """PySNMP Tutorial Common Operations Example, rewritten"""
    sysDescr = _make_object("SNMPv2-MIB", "sysDescr", 0)
    sysUpTime = _make_object("SNMPv2-MIB", "sysUpTime", 0)
    ifInOctets = _make_object("IF-MIB", "ifInOctets", 1)
    command = _make_get("192.168.0.59", "public", sysDescr, sysUpTime, ifInOctets)
    results = _run_command(command)
    _print_results(results)

def _listen_callback(
    snmp_engine,
    state_reference,
    context_engine_id,
    context_name,
    var_binds,
    callback_context,
):
    transport_domain, transport_address = snmp_engine.msgAndPduDsp.getTransportInfo(
        state_reference
    )
    print(
        f"\nNotification from {transport_address}, "
        f"SNMP Engine {context_engine_id.prettyPrint()}, "
        f"Context {context_name.prettyPrint()}"
    )
    for oid, value in var_binds:
        name = rfc1902.ObjectIdentity(oid.prettyPrint())
        name.resolveWithMib(_view_controller)
        print(f"    {name.prettyPrint()} ({oid.prettyPrint()}) = {value.prettyPrint()}")

def listen(
    address=DEFAULT_ADDRESSS,
    port=DEFAULT_PORT,
    community=DEFAULT_COMMUNITY,
    mibs=DEFAULT_MIBS,
):
    """Listen to and SNMP trap and print events."""
    # Based on pySNMP example code.
    mib_builder = builder.MibBuilder()
    compiler.addMibCompiler(mib_builder)
    mib_builder.loadModules(*mibs)
    global _view_controller
    _view_controller = view.MibViewController(mib_builder)
    loop = asyncio.get_event_loop()
    snmp_engine = engine.SnmpEngine()
    print(f"Agent is listening SNMP Trap on {address}, Port: {port}")
    if port < 1024:
        print(
            "WARNING: Port < 1024. Root priviledges or authbind required on *nix systems."
        )
    print("-" * 79)
    config.addTransport(
        snmp_engine,
        udp.domainName + (1,),
        udp.UdpTransport().openServerMode((address, port)),
    )
    config.addV1System(snmp_engine, community, community)
    ntfrcv.NotificationReceiver(snmp_engine, _listen_callback)
    print("Press CTRL-C to quit.")
    loop.run_forever()


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import atexit
import collections
import itertools
import json
import math
import os
import re
import runpy
import sys
from collections import OrderedDict
from ctypes import c_int

from . import shared_telemetry_utils as utils
from .shared_telemetry_utils import ParserError

atexit.register(ParserError.exit_func)

# Constants.
MAX_LABEL_LENGTH = 20
MAX_LABEL_COUNT = 100
MAX_KEY_COUNT = 30
MAX_KEY_LENGTH = 20
MIN_CATEGORICAL_BUCKET_COUNT = 50
CPP_IDENTIFIER_PATTERN = "^[a-z][a-z0-9_]+[a-z0-9]$"

ALWAYS_ALLOWED_KEYS = [
    "kind",
    "description",
    "operating_systems",
    "expires_in_version",
    "alert_emails",
    "keyed",
    "releaseChannelCollection",
    "bug_numbers",
    "keys",
    "record_in_processes",
    "record_into_store",
    "products",
]

BASE_DOC_URL = (
    "https://firefox-source-docs.mozilla.org/toolkit/components/" "telemetry/telemetry/"
)
HISTOGRAMS_DOC_URL = BASE_DOC_URL + "collection/histograms.html"
SCALARS_DOC_URL = BASE_DOC_URL + "collection/scalars.html"

GECKOVIEW_STREAMING_SUPPORTED_KINDS = [
    "linear",
    "exponential",
    "categorical",
]

# parse_histograms.py is used by scripts from a mozilla-central build tree
# and also by outside consumers, such as the telemetry server.  We need
# to ensure that importing things works in both contexts.  Therefore,
# unconditionally importing things that are local to the build tree, such
# as buildconfig, is a no-no.
try:
    import buildconfig

    # Need to update sys.path to be able to find usecounters.
    sys.path.append(os.path.join(buildconfig.topsrcdir, "dom/base/"))
except ImportError:
    # Must be in an out-of-tree usage scenario.  Trust that whoever is
    # running this script knows we need the usecounters module and has
    # ensured it's in our sys.path.
    pass

def linear_buckets(dmin, dmax, n_buckets):
    ret_array = [0] * n_buckets
    dmin = float(dmin)
    dmax = float(dmax)
    for i in range(1, n_buckets):
        linear_range = (dmin * (n_buckets - 1 - i) + dmax * (i - 1)) / (n_buckets - 2)
        ret_array[i] = int(linear_range + 0.5)
    return ret_array

def exponential_buckets(dmin, dmax, n_buckets):
    log_max = math.log(dmax)
    bucket_index = 2
    ret_array = [0] * n_buckets
    current = dmin
    ret_array[1] = current
    for bucket_index in range(2, n_buckets):
        log_current = math.log(current)
        log_ratio = (log_max - log_current) / (n_buckets - bucket_index)
        log_next = log_current + log_ratio
        next_value = int(math.floor(math.exp(log_next) + 0.5))
        if next_value > current:
            current = next_value
        else:
            current = current + 1
        ret_array[bucket_index] = current
    return ret_array

allowlists = None

def load_allowlist():
    global allowlists
    try:
        parsers_path = os.path.realpath(os.path.dirname(__file__))
        # The parsers live in build_scripts/parsers in the Telemetry module, while
        # the histogram-allowlists file lives in the root of the module. Account
        # for that when looking for the allowlist.
        # NOTE: if the parsers are moved, this logic will need to be updated.
        telemetry_module_path = os.path.abspath(
            os.path.join(parsers_path, os.pardir, os.pardir)
        )
        allowlist_path = os.path.join(
            telemetry_module_path, "histogram-allowlists.json"
        )
        with open(allowlist_path, "r") as f:
            try:
                allowlists = json.load(f)
                for name, allowlist in allowlists.items():
                    allowlists[name] = set(allowlist)
            except ValueError:
                ParserError("Error parsing allowlist: %s" % allowlist_path).handle_now()
    except IOError:
        allowlists = None
        ParserError("Unable to parse allowlist: %s." % allowlist_path).handle_now()

class Histogram:
    """A class for representing a histogram definition."""

    def __init__(self, name, definition, strict_type_checks=False):
        """Initialize a histogram named name with the given definition.
        definition is a dict-like object that must contain at least the keys:

         - 'kind': The kind of histogram.  Must be one of 'boolean', 'flag',
           'count', 'enumerated', 'linear', or 'exponential'.
         - 'description': A textual description of the histogram.
         - 'strict_type_checks': A boolean indicating whether to use the new, stricter type checks.
                                 The server-side still has to deal with old, oddly typed
                                 submissions, so we have to skip them there by default.
        """
        self._strict_type_checks = strict_type_checks
        self._is_use_counter = name.startswith("USE_COUNTER2_")
        if self._is_use_counter:
            definition.setdefault("record_in_processes", ["main", "content"])
            definition.setdefault("releaseChannelCollection", "opt-out")
            definition.setdefault("products", ["firefox", "fennec"])
        self.verify_attributes(name, definition)
        self._name = name
        self._description = definition["description"]
        self._kind = definition["kind"]
        self._keys = definition.get("keys", [])
        self._keyed = definition.get("keyed", False)
        self._expiration = definition.get("expires_in_version")
        self._labels = definition.get("labels", [])
        self._record_in_processes = definition.get("record_in_processes")
        self._record_into_store = definition.get("record_into_store", ["main"])
        self._products = definition.get("products")
        self._operating_systems = definition.get("operating_systems", ["all"])

        self.compute_bucket_parameters(definition)
        self.set_nsITelemetry_kind()
        self.set_dataset(definition)

    def name(self):
        """Return the name of the histogram."""
        return self._name

    def description(self):
        """Return the description of the histogram."""
        return self._description

    def kind(self):
        """Return the kind of the histogram.
        Will be one of 'boolean', 'flag', 'count', 'enumerated', 'categorical', 'linear',
        or 'exponential'."""
        return self._kind

    def expiration(self):
        """Return the expiration version of the histogram."""
        return self._expiration

    def nsITelemetry_kind(self):
        """Return the nsITelemetry constant corresponding to the kind of
        the histogram."""
        return self._nsITelemetry_kind

    def low(self):
        """Return the lower bound of the histogram."""
        return self._low

    def high(self):
        """Return the high bound of the histogram."""
        return self._high

    def n_buckets(self):
        """Return the number of buckets in the histogram."""
        return self._n_buckets

    def keyed(self):
        """Returns True if this a keyed histogram, false otherwise."""
        return self._keyed

    def keys(self):
        """Returns a list of allowed keys for keyed histogram, [] for others."""
        return self._keys

    def dataset(self):
        """Returns the dataset this histogram belongs into."""
        return self._dataset

    def labels(self):
        """Returns a list of labels for a categorical histogram, [] for others."""
        return self._labels

    def record_in_processes(self):
        """Returns a list of processes this histogram is permitted to record in."""
        return self._record_in_processes

    def record_in_processes_enum(self):
        """Get the non-empty list of flags representing the processes to record data in"""
        return [utils.process_name_to_enum(p) for p in self.record_in_processes()]

    def products(self):
        """Get the non-empty list of products to record data on"""
        return self._products

    def products_enum(self):
        """Get the non-empty list of flags representing products to record data on"""
        return [utils.product_name_to_enum(p) for p in self.products()]

    def operating_systems(self):
        """Get the list of operating systems to record data on"""
        return self._operating_systems

    def record_on_os(self, target_os):
        """Check if this probe should be recorded on the passed os."""
        os = self.operating_systems()
        if "all" in os:
            return True

        canonical_os = utils.canonical_os(target_os)

        if "unix" in os and canonical_os in utils.UNIX_LIKE_OS:
            return True

        return canonical_os in os

    def record_into_store(self):
        """Get the non-empty list of stores to record into"""
        return self._record_into_store

    def ranges(self):
        """Return an array of lower bounds for each bucket in the histogram."""
        bucket_fns = {
            "boolean": linear_buckets,
            "flag": linear_buckets,
            "count": linear_buckets,
            "enumerated": linear_buckets,
            "categorical": linear_buckets,
            "linear": linear_buckets,
            "exponential": exponential_buckets,
        }

        if self._kind not in bucket_fns:
            ParserError(
                'Unknown kind "%s" for histogram "%s".' % (self._kind, self._name)
            ).handle_later()

        fn = bucket_fns[self._kind]
        return fn(self.low(), self.high(), self.n_buckets())

    def compute_bucket_parameters(self, definition):
        bucket_fns = {
            "boolean": Histogram.boolean_flag_bucket_parameters,
            "flag": Histogram.boolean_flag_bucket_parameters,
            "count": Histogram.boolean_flag_bucket_parameters,
            "enumerated": Histogram.enumerated_bucket_parameters,
            "categorical": Histogram.categorical_bucket_parameters,
            "linear": Histogram.linear_bucket_parameters,
            "exponential": Histogram.exponential_bucket_parameters,
        }

        if self._kind not in bucket_fns:
            ParserError(
                'Unknown kind "%s" for histogram "%s".' % (self._kind, self._name)
            ).handle_later()

        fn = bucket_fns[self._kind]
        self.set_bucket_parameters(*fn(definition))

    def verify_attributes(self, name, definition):
        general_keys = ALWAYS_ALLOWED_KEYS + ["low", "high", "n_buckets"]

        table = {
            "boolean": ALWAYS_ALLOWED_KEYS,
            "flag": ALWAYS_ALLOWED_KEYS,
            "count": ALWAYS_ALLOWED_KEYS,
            "enumerated": ALWAYS_ALLOWED_KEYS + ["n_values"],
            "categorical": ALWAYS_ALLOWED_KEYS + ["labels", "n_values"],
            "linear": general_keys,
            "exponential": general_keys,
        }
        # We removed extended_statistics_ok on the client, but the server-side,
        # where _strict_type_checks==False, has to deal with historical data.
        if not self._strict_type_checks:
            table["exponential"].append("extended_statistics_ok")

        kind = definition["kind"]
        if kind not in table:
            ParserError(
                'Unknown kind "%s" for histogram "%s".' % (kind, name)
            ).handle_later()
        allowed_keys = table[kind]

        self.check_name(name)
        self.check_keys(name, definition, allowed_keys)
        self.check_keys_field(name, definition)
        self.check_field_types(name, definition)
        self.check_allowlisted_kind(name, definition)
        self.check_allowlistable_fields(name, definition)
        self.check_expiration(name, definition)
        self.check_label_values(name, definition)
        self.check_record_in_processes(name, definition)
        self.check_products(name, definition)
        self.check_operating_systems(name, definition)
        self.check_record_into_store(name, definition)

    def check_name(self, name):
        if "#" in name:
            ParserError(
                'Error for histogram name "%s": "#" is not allowed.' % (name)
            ).handle_later()

        # Avoid C++ identifier conflicts between histogram enums and label enum names.
        if name.startswith("LABELS_"):
            ParserError(
                'Error for histogram name "%s":  can not start with "LABELS_".' % (name)
            ).handle_later()

        # To make it easier to generate C++ identifiers from this etc., we restrict
        # the histogram names to a strict pattern.
        # We skip this on the server to avoid failures with old Histogram.json revisions.
        if self._strict_type_checks:
            if not re.match(CPP_IDENTIFIER_PATTERN, name, re.IGNORECASE):
                ParserError(
                    'Error for histogram name "%s": name does not conform to "%s"'
                    % (name, CPP_IDENTIFIER_PATTERN)
                ).handle_later()

    def check_expiration(self, name, definition):
        field = "expires_in_version"
        expiration = definition.get(field)

        if not expiration:
            return

        # We forbid new probes from using "expires_in_version" : "default" field/value pair.
        # Old ones that use this are added to the allowlist.
        if (
            expiration == "default"
            and allowlists is not None
            and name not in allowlists["expiry_default"]
        ):
            ParserError(
                'New histogram "%s" cannot have "default" %s value.' % (name, field)
            ).handle_later()

        # Historical editions of Histograms.json can have the deprecated
        # expiration format 'N.Na1'. Fortunately, those scripts set
        # self._strict_type_checks to false.
        if (
            expiration != "default"
            and not utils.validate_expiration_version(expiration)
            and self._strict_type_checks
        ):
            ParserError(
                (
                    "Error for histogram {} - invalid {}: {}."
                    "\nSee: {}#expires-in-version"
                ).format(name, field, expiration, HISTOGRAMS_DOC_URL)
            ).handle_later()

        expiration = utils.add_expiration_postfix(expiration)

        definition[field] = expiration

    def check_label_values(self, name, definition):
        labels = definition.get("labels")
        if not labels:
            return

        invalid = filter(lambda l: len(l) > MAX_LABEL_LENGTH, labels)
        if len(list(invalid)) > 0:
            ParserError(
                'Label values for "%s" exceed length limit of %d: %s'
                % (name, MAX_LABEL_LENGTH, ", ".join(invalid))
            ).handle_later()

        if len(labels) > MAX_LABEL_COUNT:
            ParserError(
                'Label count for "%s" exceeds limit of %d' % (name, MAX_LABEL_COUNT)
            ).handle_now()

        # To make it easier to generate C++ identifiers from this etc., we restrict
        # the label values to a strict pattern.
        invalid = filter(
            lambda l: not re.match(CPP_IDENTIFIER_PATTERN, l, re.IGNORECASE), labels
        )
        if len(list(invalid)) > 0:
            ParserError(
                'Label values for %s are not matching pattern "%s": %s'
                % (name, CPP_IDENTIFIER_PATTERN, ", ".join(invalid))
            ).handle_later()

    def check_record_in_processes(self, name, definition):
        if not self._strict_type_checks:
            return

        field = "record_in_processes"
        rip = definition.get(field)

        DOC_URL = HISTOGRAMS_DOC_URL + "#record-in-processes"

        if not rip:
            ParserError(
                'Histogram "%s" must have a "%s" field:\n%s' % (name, field, DOC_URL)
            ).handle_later()

        for process in rip:
            if not utils.is_valid_process_name(process):
                ParserError(
                    'Histogram "%s" has unknown process "%s" in %s.\n%s'
                    % (name, process, field, DOC_URL)
                ).handle_later()

    def check_products(self, name, definition):
        if not self._strict_type_checks:
            return

        field = "products"
        products = definition.get(field)

        DOC_URL = HISTOGRAMS_DOC_URL + "#products"

        if not products:
            ParserError(
                'Histogram "%s" must have a "%s" field:\n%s' % (name, field, DOC_URL)
            ).handle_now()

        for product in products:
            if not utils.is_valid_product(product):
                ParserError(
                    'Histogram "%s" has unknown product "%s" in %s.\n%s'
                    % (name, product, field, DOC_URL)
                ).handle_later()
            if utils.is_geckoview_streaming_product(product):
                kind = definition.get("kind")
                if kind not in GECKOVIEW_STREAMING_SUPPORTED_KINDS:
                    ParserError(
                        (
                            'Histogram "%s" is of kind "%s" which is unsupported for '
                            'product "%s".'
                        )
                        % (name, kind, product)
                    ).handle_later()
                keyed = definition.get("keyed")
                if keyed:
                    ParserError(
                        'Keyed histograms like "%s" are unsupported for product "%s"'
                        % (name, product)
                    ).handle_later()

    def check_operating_systems(self, name, definition):
        if not self._strict_type_checks:
            return

        field = "operating_systems"
        operating_systems = definition.get(field)

        DOC_URL = HISTOGRAMS_DOC_URL + "#operating-systems"

        if not operating_systems:
            # operating_systems is optional
            return

        for operating_system in operating_systems:
            if not utils.is_valid_os(operating_system):
                ParserError(
                    'Histogram "%s" has unknown operating system "%s" in %s.\n%s'
                    % (name, operating_system, field, DOC_URL)
                ).handle_later()

    def check_record_into_store(self, name, definition):
        if not self._strict_type_checks:
            return

        field = "record_into_store"
        DOC_URL = HISTOGRAMS_DOC_URL + "#record-into-store"

        if field not in definition:
            # record_into_store is optional
            return

        record_into_store = definition.get(field)
        # record_into_store should not be empty
        if not record_into_store:
            ParserError(
                'Histogram "%s" has empty list of stores, which is not allowed.\n%s'
                % (name, DOC_URL)
            ).handle_later()

    def check_keys_field(self, name, definition):
        keys = definition.get("keys")
        if not self._strict_type_checks or keys is None:
            return

        if not definition.get("keyed", False):
            raise ValueError(
                "'keys' field is not valid for %s; only allowed for keyed histograms."
                % (name)
            )

        if len(keys) == 0:
            raise ValueError("The key list for %s cannot be empty" % (name))

        if len(keys) > MAX_KEY_COUNT:
            raise ValueError(
                "Label count for %s exceeds limit of %d" % (name, MAX_KEY_COUNT)
            )

        invalid = filter(lambda k: len(k) > MAX_KEY_LENGTH, keys)
        if len(list(invalid)) > 0:
            raise ValueError(
                '"keys" values for %s are exceeding length "%d": %s'
                % (name, MAX_KEY_LENGTH, ", ".join(invalid))
            )

    def check_allowlisted_kind(self, name, definition):
        # We don't need to run any of these checks on the server.
        if not self._strict_type_checks or allowlists is None:
            return

        # Disallow "flag" and "count" histograms on desktop, suggest to use
        # scalars instead. Allow using these histograms on Android, as we
        # don't support scalars there yet.
        hist_kind = definition.get("kind")
        android_target = "android" in definition.get("operating_systems", [])

        if (
            not android_target
            and hist_kind in ["flag", "count"]
            and name not in allowlists["kind"]
        ):
            ParserError(
                (
                    'Unsupported kind "%s" for histogram "%s":\n'
                    'New "%s" histograms are not supported on Desktop, you should'
                    " use scalars instead:\n"
                    "%s\n"
                    "Are you trying to add a histogram on Android?"
                    ' Add "operating_systems": ["android"] to your histogram definition.'
                )
                % (hist_kind, name, hist_kind, SCALARS_DOC_URL)
            ).handle_now()

    # Check for the presence of fields that old histograms are allowlisted for.
    def check_allowlistable_fields(self, name, definition):
        # Use counters don't have any mechanism to add the fields checked here,
        # so skip the check for them.
        # We also don't need to run any of these checks on the server.
        if self._is_use_counter or not self._strict_type_checks:
            return

        # In the pipeline we don't have allowlists available.
        if allowlists is None:
            return

        for field in ["alert_emails", "bug_numbers"]:
            if field not in definition and name not in allowlists[field]:
                ParserError(
                    'New histogram "%s" must have a "%s" field.' % (name, field)
                ).handle_later()
            if field in definition and name in allowlists[field]:
                msg = (
                    'Histogram "%s" should be removed from the allowlist for "%s" in '
                    "histogram-allowlists.json."
                )
                ParserError(msg % (name, field)).handle_later()

    def check_field_types(self, name, definition):
        # Define expected types for the histogram properties.
        type_checked_fields = {
            "n_buckets": int,
            "n_values": int,
            "low": int,
            "high": int,
            "keyed": bool,
            "expires_in_version": str,
            "kind": str,
            "description": str,
            "releaseChannelCollection": str,
        }

        # For list fields we check the items types.
        type_checked_list_fields = {
            "bug_numbers": int,
            "alert_emails": str,
            "labels": str,
            "record_in_processes": str,
            "keys": str,
            "products": str,
            "operating_systems": str,
            "record_into_store": str,
        }

        # For the server-side, where _strict_type_checks==False, we want to
        # skip the stricter type checks for these fields for dealing with
        # historical data.
        coerce_fields = ["low", "high", "n_values", "n_buckets"]
        if not self._strict_type_checks:
            # This handles some old non-numeric expressions.
            EXPRESSIONS = {
                "JS::GCReason::NUM_TELEMETRY_REASONS": 101,
                "mozilla::StartupTimeline::MAX_EVENT_ID": 12,
            }

            def try_to_coerce_to_number(v):
                if v in EXPRESSIONS:
                    return EXPRESSIONS[v]
                try:
                    return eval(v, {})
                except Exception:
                    return v

            for key in [k for k in coerce_fields if k in definition]:
                definition[key] = try_to_coerce_to_number(definition[key])
            # This handles old "keyed":"true" definitions (bug 1271986).
            if definition.get("keyed", None) == "true":
                definition["keyed"] = True

        def nice_type_name(t):
            if t is str:
                return "string"
            return t.__name__

        for key, key_type in type_checked_fields.items():
            if key not in definition:
                continue
            if not isinstance(definition[key], key_type):
                ParserError(
                    'Value for key "{0}" in histogram "{1}" should be {2}.'.format(
                        key, name, nice_type_name(key_type)
                    )
                ).handle_later()

        # Make sure the max range is lower than or equal to INT_MAX
        if "high" in definition and not c_int(definition["high"]).value > 0:
            ParserError(
                'Value for high in histogram "{0}" should be lower or equal to INT_MAX.'.format(
                    nice_type_name(c_int)
                )
            ).handle_later()

        for key, key_type in type_checked_list_fields.items():
            if key not in definition:
                continue
            if not all(isinstance(x, key_type) for x in definition[key]):
                ParserError(
                    'All values for list "{0}" in histogram "{1}" should be of type'
                    " {2}.".format(key, name, nice_type_name(key_type))
                ).handle_later()

    def check_keys(self, name, definition, allowed_keys):
        if not self._strict_type_checks:
            return
        for key in iter(definition.keys()):
            if key not in allowed_keys:
                ParserError(
                    'Key "%s" is not allowed for histogram "%s".' % (key, name)
                ).handle_later()

    def set_bucket_parameters(self, low, high, n_buckets):
        self._low = low
        self._high = high
        self._n_buckets = n_buckets
        max_n_buckets = 101 if self._kind in ["enumerated", "categorical"] else 100
        if (
            allowlists is not None
            and self._n_buckets > max_n_buckets
            and type(self._n_buckets) is int
        ):
            if self._name not in allowlists["n_buckets"]:
                ParserError(
                    'New histogram "%s" is not permitted to have more than 100 buckets.\n'
                    "Histograms with large numbers of buckets use disproportionately high"
                    " amounts of resources. Contact a Telemetry peer (e.g. in #telemetry)"
                    " if you think an exception ought to be made:\n"
                    "https://wiki.mozilla.org/Modules/Toolkit#Telemetry" % self._name
                ).handle_later()

    @staticmethod
    def boolean_flag_bucket_parameters(definition):
        return (1, 2, 3)

    @staticmethod
    def linear_bucket_parameters(definition):
        return (definition.get("low", 1), definition["high"], definition["n_buckets"])

    @staticmethod
    def enumerated_bucket_parameters(definition):
        n_values = definition["n_values"]
        return (1, n_values, n_values + 1)

    @staticmethod
    def categorical_bucket_parameters(definition):
        # Categorical histograms default to 50 buckets to make working with them easier.
        # Otherwise when adding labels later we run into problems with the pipeline not
        # supporting bucket changes.
        # This can be overridden using the n_values field.
        n_values = max(
            len(definition["labels"]),
            definition.get("n_values", 0),
            MIN_CATEGORICAL_BUCKET_COUNT,
        )
        return (1, n_values, n_values + 1)

    @staticmethod
    def exponential_bucket_parameters(definition):
        return (definition.get("low", 1), definition["high"], definition["n_buckets"])

    def set_nsITelemetry_kind(self):
        # Pick a Telemetry implementation type.
        types = {
            "boolean": "BOOLEAN",
            "flag": "FLAG",
            "count": "COUNT",
            "enumerated": "LINEAR",
            "categorical": "CATEGORICAL",
            "linear": "LINEAR",
            "exponential": "EXPONENTIAL",
        }

        if self._kind not in types:
            ParserError(
                'Unknown kind "%s" for histogram "%s".' % (self._kind, self._name)
            ).handle_later()

        self._nsITelemetry_kind = "nsITelemetry::HISTOGRAM_%s" % types[self._kind]

    def set_dataset(self, definition):
        datasets = {
            "opt-in": "DATASET_PRERELEASE_CHANNELS",
            "opt-out": "DATASET_ALL_CHANNELS",
        }

        value = definition.get("releaseChannelCollection", "opt-in")
        if value not in datasets:
            ParserError(
                "Unknown value for releaseChannelCollection"
                ' policy for histogram "%s".' % self._name
            ).handle_later()

        self._dataset = "nsITelemetry::" + datasets[value]

# This hook function loads the histograms into an OrderedDict.
# It will raise a ParserError if duplicate keys are found.
def load_histograms_into_dict(ordered_pairs, strict_type_checks):
    d = collections.OrderedDict()
    for key, value in ordered_pairs:
        if strict_type_checks and key in d:
            ParserError(
                "Found duplicate key in Histograms file: %s" % key
            ).handle_later()
        d[key] = value
    return d

# We support generating histograms from multiple different input files, not
# just Histograms.json.  For each file's basename, we have a specific
# routine to parse that file, and return a dictionary mapping histogram
# names to histogram parameters.
def from_json(filename, strict_type_checks):
    with open(filename, "r") as f:
        try:

            def hook(ps):
                return load_histograms_into_dict(ps, strict_type_checks)

            histograms = json.load(f, object_pairs_hook=hook)
        except ValueError as e:
            ParserError(
                "error parsing histograms in %s: %s" % (filename, e)
            ).handle_now()
    return histograms

def from_UseCounters_conf(filename, strict_type_checks):
    return usecounters.generate_histograms(filename)

def from_UseCountersWorker_conf(filename, strict_type_checks):
    return usecounters.generate_histograms(filename, True)

def from_nsDeprecatedOperationList(filename, strict_type_checks):
    operation_regex = re.compile("^DEPRECATED_OPERATION\\(([^)]+)\\)")
    histograms = collections.OrderedDict()

    with open(filename, "r") as f:
        for line in f:
            match = operation_regex.search(line)
            if not match:
                continue

            op = match.group(1)

            def add_counter(context):
                name = "USE_COUNTER2_DEPRECATED_%s_%s" % (op, context.upper())
                histograms[name] = {
                    "expires_in_version": "never",
                    "kind": "boolean",
                    "description": "Whether a %s used %s" % (context, op),
                }

            add_counter("document")
            add_counter("page")

    return histograms

def to_camel_case(property_name):
    return re.sub(
        "(^|_|-)([a-z0-9])",
        lambda m: m.group(2).upper(),
        property_name.strip("_").strip("-"),
    )

def add_css_property_counters(histograms, property_name):
    def add_counter(context):
        name = "USE_COUNTER2_CSS_PROPERTY_%s_%s" % (
            to_camel_case(property_name),
            context.upper(),
        )
        histograms[name] = {
            "expires_in_version": "never",
            "kind": "boolean",
            "description": "Whether a %s used the CSS property %s"
            % (context, property_name),
        }

    add_counter("document")
    add_counter("page")

def from_ServoCSSPropList(filename, strict_type_checks):
    histograms = collections.OrderedDict()
    properties = runpy.run_path(filename)["data"]
    for prop in properties:
        add_css_property_counters(histograms, prop.name)
    return histograms

def from_counted_unknown_properties(filename, strict_type_checks):
    histograms = collections.OrderedDict()
    properties = runpy.run_path(filename)["COUNTED_UNKNOWN_PROPERTIES"]

    # NOTE(emilio): Unlike ServoCSSProperties, `prop` here is just the property
    # name.
    #
    # We use the same naming as CSS properties so that we don't get
    # discontinuity when we implement or prototype them.
    for prop in properties:
        add_css_property_counters(histograms, prop)
    return histograms

# This is only used for probe-scraper.
def from_properties_db(filename, strict_type_checks):
    histograms = collections.OrderedDict()
    with open(filename, "r") as f:
        in_css_properties = False

        for line in f:
            if not in_css_properties:
                if line.startswith("exports.CSS_PROPERTIES = {"):
                    in_css_properties = True
                continue

            if line.startswith("};"):
                break

            if not line.startswith('  "'):
                continue

            name = line.split('"')[1]
            add_css_property_counters(histograms, name)
    return histograms

FILENAME_PARSERS = [
    (lambda x: from_json if x.endswith(".json") else None),
    (
        lambda x: from_nsDeprecatedOperationList
        if x == "nsDeprecatedOperationList.h"
        else None
    ),
    (lambda x: from_ServoCSSPropList if x == "ServoCSSPropList.py" else None),
    (
        lambda x: from_counted_unknown_properties
        if x == "counted_unknown_properties.py"
        else None
    ),
    (lambda x: from_properties_db if x == "properties-db.js" else None),
]

# Similarly to the dance above with buildconfig, usecounters may not be
# available, so handle that gracefully.
try:
    import usecounters

    FILENAME_PARSERS.append(
        lambda x: from_UseCounters_conf if x == "UseCounters.conf" else None
    )
    FILENAME_PARSERS.append(
        lambda x: from_UseCountersWorker_conf if x == "UseCountersWorker.conf" else None
    )
except ImportError:
    pass

def from_files(filenames, strict_type_checks=True):
    """Return an iterator that provides a sequence of Histograms for
    the histograms defined in filenames.
    """
    if strict_type_checks:
        load_allowlist()

    all_histograms = OrderedDict()
    for filename in filenames:
        parser = None
        for checkFn in FILENAME_PARSERS:
            parser = checkFn(os.path.basename(filename))
            if parser is not None:
                break

        if parser is None:
            ParserError("Don't know how to parse %s." % filename).handle_now()

        histograms = parser(filename, strict_type_checks)

        # OrderedDicts are important, because then the iteration order over
        # the parsed histograms is stable, which makes the insertion into
        # all_histograms stable, which makes ordering in generated files
        # stable, which makes builds more deterministic.
        if not isinstance(histograms, OrderedDict):
            ParserError("Histogram parser did not provide an OrderedDict.").handle_now()

        for (name, definition) in histograms.items():
            if name in all_histograms:
                ParserError('Duplicate histogram name "%s".' % name).handle_later()
            all_histograms[name] = definition

    def check_continuity(iterable, filter_function, name):
        indices = list(filter(filter_function, enumerate(iter(iterable.keys()))))
        if indices:
            lower_bound = indices[0][0]
            upper_bound = indices[-1][0]
            n_counters = upper_bound - lower_bound + 1
            if n_counters != len(indices):
                ParserError(
                    "Histograms %s must be defined in a contiguous block." % name
                ).handle_later()

    # We require that all USE_COUNTER2_*_WORKER histograms be defined in a contiguous
    # block.
    check_continuity(
        all_histograms,
        lambda x: x[1].startswith("USE_COUNTER2_") and x[1].endswith("_WORKER"),
        "use counter worker",
    )
    # And all other USE_COUNTER2_* histograms be defined in a contiguous
    # block.
    check_continuity(
        all_histograms,
        lambda x: x[1].startswith("USE_COUNTER2_") and not x[1].endswith("_WORKER"),
        "use counter",
    )

    # Check that histograms that were removed from Histograms.json etc.
    # are also removed from the allowlists.
    if allowlists is not None:
        all_allowlist_entries = itertools.chain.from_iterable(iter(allowlists.values()))
        orphaned = set(all_allowlist_entries) - set(all_histograms.keys())
        if len(orphaned) > 0:
            msg = (
                "The following entries are orphaned and should be removed from "
                "histogram-allowlists.json:\n%s"
            )
            ParserError(msg % (", ".join(sorted(orphaned)))).handle_later()

    for (name, definition) in all_histograms.items():
        yield Histogram(name, definition, strict_type_checks=strict_type_checks)


from __future__ import absolute_import, division, print_function

__metaclass__ = type

from ansible_collections.community.isva.plugins.module_utils.isva_utils import parse_fail_message, identical_files
from ansible_collections.community.isva.plugins.module_utils.common import ISVAModuleError
from ansible.module_utils.connection import Connection

import json
import logging
import os
import base64
import shutil

uri = '/isam/downloads'

logger = logging.getLogger(__name__)

MAP_API_ATTRIBUTES = {}

def from_api(source):
    data = {}
    for _, value in MAP_API_ATTRIBUTES.items():
        if value in source and source[value] is not None:
            data = source[value]

    return data

def fetch_file_downloads(module, path=None, recursive=True):
    """ This function fetch the file downloads from the appliance for a given path

    Returns:
        _type_: _description_
    """
    if path:
        target_uri = '{}/{}?recursive={}'.format(uri, path, recursive)
    else:
        target_uri = '{}?recursive={}'.format(uri, recursive)

    connection = Connection(module._socket_path)
    response = connection.send_request(path=target_uri)

    if response['code'] != 200:
        raise ISVAModuleError(parse_fail_message(response['code'], response['contents']))

    return response['contents']

def download_file_downloads(module, path, dest, remote_files={}):
    """This function will download the requested shared volume file.

    """
    target_uri = '{}/{}'.format(uri, path)
    file_path = path.split('/')

    curr_d = remote_files
    for d in file_path[1:]:  # Starting at 1 to avoid the first empty string
        curr_d = curr_d.get(d)
        if not curr_d:
            raise ISVAModuleError('The requested file does not exist {}'.format(path))

    tmp_dest = '{}{}'.format(module.tmpdir, file_path[-1])
    # Downloading file, we've no way to find out remotely if they are identical, so adding it at a temporary location.
    connection = Connection(module._socket_path)
    response = connection.download_file(path=target_uri, dest=tmp_dest)

    if not response:
        raise ISVAModuleError('Couldn\'t download the file {}/{}'.format(volume, path))

    if identical_files(module, tmp_dest, dest):  # No change to report
        return False

    shutil.copyfile(tmp_dest, dest)

    return True


import copy
import numpy as np
from sklearn.utils import shuffle
from tensorflow import keras

from shapley.loader import Loader

class MNIST(Loader):
    def __init__(self, num_train, one_hot=True, shuffle=True, by_label=False):
        self.name = 'mnist'
        self.num_train = num_train
        self.num_test = num_train // 10
        self.x_train, self.y_train, self.x_test, self.y_test = self.load_data(one_hot, by_label)
        if shuffle: self.shuffle_data()
        self.x_train = self.x_train[:self.num_train]
        self.y_train = self.y_train[:self.num_train]
        self.x_test = self.x_test[:self.num_test]
        self.y_test = self.y_test[:self.num_test]

    def load_data(self, one_hot, by_label):
        mnist = keras.datasets.mnist
        (x_train, y_train), (x_test, y_test) = mnist.load_data()
        x_train = np.reshape(x_train, [-1, 28 * 28])
        x_train = x_train.astype(np.float32) / 255
        x_test = np.reshape(x_test, [-1, 28 * 28])
        x_test = x_test.astype(np.float32) / 255

        if by_label:
            ind_train = np.argsort(y_train)
            ind_test = np.argsort(y_test)
            x_train, y_train = x_train[ind_train], y_train[ind_train]
            x_test, y_test = x_test[ind_test], y_test[ind_test]

        if one_hot:
            # convert to one-hot labels
            y_train = keras.utils.to_categorical(y_train)
            y_test = keras.utils.to_categorical(y_test)

        return x_train, y_train, x_test, y_test

    def shuffle_data(self):
        ind = np.random.permutation(len(self.x_train))
        self.x_train, self.y_train = self.x_train[ind], self.y_train[ind]
        ind = np.random.permutation(len(self.x_test))
        self.x_test, self.y_test = self.x_test[ind], self.y_test[ind]

    def prepare_data(self):
        return self.x_train, self.y_train, self.x_test, self.y_test


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision
import pytorch_lightning as pl

from argparse import ArgumentParser, Namespace
from collections import OrderedDict

from dataloaders.dataloader_self import Multiview_Dataset
from utils.radam import RAdam

import numpy as np
eps = torch.tensor(np.finfo(float).eps).float()

class Discriminator(nn.Module):

    def __init__(self, tile_size=(64, 64, 64), in_channels=1, **kwargs):
        super(Discriminator, self).__init__()

        self.tile_size = tile_size
        self.in_channels = in_channels
        self.out_size = tuple([int(p/2**4) for p in tile_size])

        # Define layer instances
        self.conv1 = nn.Conv3d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=2, padding=1)
        self.norm2 = nn.InstanceNorm3d(num_features=64)
        self.leaky1 = nn.LeakyReLU(negative_slope=0.2)

        self.conv2 = nn.Conv3d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1)
        self.norm2 = nn.InstanceNorm3d(num_features=128)
        self.leaky2 = nn.LeakyReLU(negative_slope=0.2)

        self.conv3 = nn.Conv3d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)
        self.norm3 = nn.InstanceNorm3d(num_features=256)
        self.leaky3 = nn.LeakyReLU(negative_slope=0.2)

        self.conv4 = nn.Conv3d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1)
        self.norm4 = nn.InstanceNorm3d(num_features=512)
        self.leaky4 = nn.LeakyReLU(negative_slope=0.2)

        self.conv6 = nn.Conv3d(in_channels=512, out_channels=1, kernel_size=3, stride=1, padding=1)
        self.sig6 = nn.Sigmoid()

    def forward(self, img):

        out = self.conv1(img)
        out = self.leaky1(out)

        out = self.conv2(out)
        out = self.norm2(out)
        out = self.leaky2(out)

        out = self.conv3(out)
        out = self.norm3(out)
        out = self.leaky3(out)

        out = self.conv4(out)
        out = self.norm4(out)
        out = self.leaky4(out)

        out = self.conv6(out)
        out = self.sig6(out)

        return out

class Generator(nn.Module):
    """
    Implementation of the 3D U-Net architecture.
    """

    def __init__(self, in_channels=4, out_channels=1, feat_channels=64, num_views=4,
        out_activation='sigmoid', norm_method='instance', data_augmentation=True, **kwargs):
        super(Generator, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.feat_channels = feat_channels
        self.num_views = num_views
        self.out_activation = out_activation # relu | sigmoid | tanh | hardtanh | none
        self.norm_method = norm_method # instance | batch | none

        if self.norm_method == 'instance':
            self.norm = nn.InstanceNorm3d
        elif self.norm_method == 'batch':
            self.norm = nn.BatchNorm3d
        elif self.norm_method == 'none':
            self.norm = nn.Identity
        else:
            raise ValueError('Unknown normalization method "{0}". Choose from "instance|batch|none".'.format(self.norm_method))

        # The forward generator (point spread functions)
        self.PSF_conv = []
        self.data_augmentation = data_augmentation  # If true, the forward generator has to be re-assigned for a new batch
        self.kernels_initilized = False  # Whether the forward generator has been initialized

        # Define layer instances of the inverse generator (deconvolution network)
        self.c1 = nn.Sequential(
            nn.Conv3d(in_channels=in_channels, out_channels=feat_channels, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=feat_channels, out_channels=feat_channels, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels),
            nn.ReLU(inplace=True)
            )
        self.d1 = nn.Sequential(nn.MaxPool3d(kernel_size=2, stride=2, padding=0))

        self.c2 = nn.Sequential(
            nn.Conv3d(in_channels=feat_channels, out_channels=feat_channels*2, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels*2),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=feat_channels*2, out_channels=feat_channels*2, kernel_size=3, padding=1),
            self.norm(num_features=feat_channels*2),
            nn.ReLU(inplace=True)
            )
        self.d2 = nn.Sequential(nn.MaxPool3d(kernel_size=2, stride=2, padding=0))

        self.c3 = nn.Sequential(
            nn.Conv3d(in_channels=feat_channels*2, out_channels=feat_channels*4, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels*4),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=feat_channels*4, out_channels=feat_channels*4, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels*4),
            nn.ReLU(inplace=True)
            )

        self.u1 = nn.Sequential(
            nn.ConvTranspose3d(in_channels=feat_channels*4, out_channels=feat_channels*2, kernel_size=4, stride=2, padding=1, output_padding=0),
            self.norm(num_features=feat_channels*4),
            nn.ReLU(inplace=True)
            )
        self.c4 = nn.Sequential(
            nn.Conv3d(in_channels=feat_channels*4, out_channels=feat_channels*2, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels*2),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=feat_channels*2, out_channels=feat_channels*2, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels*2),
            nn.ReLU(inplace=True)
            )

        self.u2 = nn.Sequential(
            nn.ConvTranspose3d(in_channels=feat_channels*2, out_channels=feat_channels, kernel_size=4, stride=2, padding=1, output_padding=0),
            self.norm(num_features=feat_channels*2),
            nn.ReLU(inplace=True)
            )
        self.c5 = nn.Sequential(
            nn.Conv3d(in_channels=feat_channels*2, out_channels=feat_channels, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=feat_channels, out_channels=feat_channels, kernel_size=3, stride=1, padding=1),
            self.norm(num_features=feat_channels),
            nn.ReLU(inplace=True)
            )

        self.out = nn.Sequential(
            nn.Conv3d(in_channels=feat_channels, out_channels=out_channels, kernel_size=1, stride=1)
            )

        if self.out_activation == 'relu':
            self.out_fcn = nn.ReLU()
        elif self.out_activation == 'sigmoid':
            self.out_fcn = nn.Sigmoid()
        elif self.out_activation == 'tanh':
            self.out_fcn = nn.Tanh()
        elif self.out_activation == 'leakyrelu':
            self.out_fcn = nn.LeakyReLU(negative_slope=0.1)
        elif self.out_activation == 'none':
            self.out_fcn = None
        else:
            raise ValueError('Unknown output activation "{0}". Choose from "relu|sigmoid|tanh|hardtanh|none".'.format(self.out_activation))

    def set_kernel_weights(self, PSF):
        '''
        Initialize convolutional kernels with the given PSF
        '''

        for i in range(self.num_views):
            psf_size = tuple(PSF[0, i, ...].size())
            len_padding = tuple([(s - 1) // 2 for s in psf_size])

            if not self.kernels_initilized:
                self.PSF_conv.append(nn.Conv3d(in_channels=self.out_channels, out_channels=self.out_channels, \
                    kernel_size=psf_size, stride=1, padding=len_padding).cuda())

            normalized = (PSF[0, i, ...] / PSF[0, i, ...].sum()).unsqueeze(0).unsqueeze(0)

            self.PSF_conv[i].weight = nn.Parameter(normalized, requires_grad=False)
            self.PSF_conv[i].bias = nn.Parameter(torch.zeros_like(self.PSF_conv[i].bias), requires_grad=False)

        self.kernels_initilized = True

    def forward(self, img, PSF, direction):

        # Blurry-->sharp-->blurry
        if direction == 'BSB':

            c1 = self.c1(img)
            d1 = self.d1(c1)
            # print('c1 ', c1.size())
            # print('d1 ', d1.size())

            c2 = self.c2(d1)
            d2 = self.d2(c2)
            # print('c2 ', c2.size())
            # print('d2 ', d2.size())

            c3 = self.c3(d2)
            # print('c3 ', c3.size())

            u1 = self.u1(c3)
            c4 = self.c4(torch.cat((u1, c2), dim=1))
            # print('u1 ', u1.size())
            # print('c4 ', c4.size())

            u2 = self.u2(c4)
            c5 = self.c5(torch.cat((u2, c1), dim=1))
            # print('u2 ', u2.size())
            # print('c5 ', c5.size())

            if self.out_fcn is None:
                sharp_pred = self.out(c5)
            else:
                sharp_pred = self.out_fcn(self.out(c5))

            if PSF is None:
                blurred_with_PSF = None
            else:
                if self.data_augmentation or not self.kernels_initilized:
                    self.set_kernel_weights(PSF)

                blurred_with_PSF = []
                for i in range(self.num_views):
                    blurred_with_PSF.append(self.PSF_conv[i](sharp_pred))
                blurred_with_PSF = torch.cat(blurred_with_PSF, dim=1)

        # Sharp-->blurry-->sharp
        elif direction == 'SBS':

            if PSF is None:
                blurred_with_PSF = None
            else:
                if self.data_augmentation or not self.kernels_initilized:
                    self.set_kernel_weights(PSF)

                blurred_with_PSF = []
                for i in range(self.num_views):
                    blurred_with_PSF.append(self.PSF_conv[i](img))
                blurred_with_PSF = torch.cat(blurred_with_PSF, dim=1)

            c1 = self.c1(blurred_with_PSF)
            d1 = self.d1(c1)
            # print('c1 ', c1.size())
            # print('d1 ', d1.size())

            c2 = self.c2(d1)
            d2 = self.d2(c2)
            # print('c2 ', c2.size())
            # print('d2 ', d2.size())

            c3 = self.c3(d2)
            # print('c3 ', c3.size())

            u1 = self.u1(c3)
            c4 = self.c4(torch.cat((u1, c2), dim=1))
            # print('u1 ', u1.size())
            # print('c4 ', c4.size())

            u2 = self.u2(c4)
            c5 = self.c5(torch.cat((u2, c1), dim=1))
            # print('u2 ', u2.size())
            # print('c5 ', c5.size())

            if self.out_fcn is None:
                sharp_pred = self.out(c5)
            else:
                sharp_pred = self.out_fcn(self.out(c5))

        return sharp_pred, blurred_with_PSF

class cycleGAN(pl.LightningModule):

    def __init__(self, hparams):
        super().__init__()
        if type(hparams) is dict:
            hparams = Namespace(**hparams)
        self.hparams = hparams

        # Generator
        self.generator = Generator(in_channels=hparams.in_channels, out_channels=hparams.out_channels,
            feat_channels=hparams.feat_channels, num_views=hparams.num_views, out_activation=hparams.out_activation,
            norm_method=hparams.norm_method, data_augmentation=hparams.data_augmentation)

        # Discriminator with full-size tile
        self.discriminator_full = Discriminator(tile_size=hparams.tile_size, in_channels=hparams.out_channels)

        # Discriminator with half-size tile
        half_tile_size = (dim // 2 for dim in hparams.tile_size)
        self.discriminator_half = Discriminator(tile_size=half_tile_size, in_channels=hparams.out_channels)

        # Cache for generated images
        self.last_views = None
        self.last_BSB_sharp_pred = None
        self.last_SBS_sharp_pred = None
        self.last_gt = None
        self.slicing_half = None

        # Cycle loss function
        if hparams.cycle_loss_function == 'l1':
            self.cycle_loss_func = F.l1_loss
        else:
            self.cycle_loss_func = F.mse_loss

    def forward(self, img, PSF=None, direction='BSB'):
        return self.generator(img, PSF, direction)

    def compute_cycle_loss(self, img_before_cycle, img_after_cycle):
        return self.cycle_loss_func(img_before_cycle, img_after_cycle, reduction='mean')

    def compute_LSGAN_loss(self, pred, mask):
        return F.mse_loss(pred, mask, reduction='mean')

    def compute_gradient_loss(self, img):
        """
        3D adaption of numpy.gradient()
        """

        dz = torch.zeros_like(img)
        dy = torch.zeros_like(img)
        dx = torch.zeros_like(img)
        out = [dz, dy, dx]
        N = img.ndimension()
        s_dim = N - 2  # Number of spatial dimensions

        # Initialize slice objects with [:, ..., :]
        slice1 = [slice(None)]*N
        slice2 = [slice(None)]*N
        slice3 = [slice(None)]*N
        slice4 = [slice(None)]*N

        for i, axis in enumerate(range(-s_dim, 0)):

            # 2nd order difference in the interior
            slice1[axis] = slice(1, -1)
            slice2[axis] = slice(None, -2)
            slice3[axis] = slice(1, -1)
            slice4[axis] = slice(2, None)

            out[i][tuple(slice1)] = (img[tuple(slice4)] - img[tuple(slice2)]) / 2

            # 1st order difference on edges
            slice1[axis] = 0
            slice2[axis] = 1
            slice3[axis] = 0

            # 1D equivalent -- out[0] = (f[1] - f[0]) / (x[1] - x[0])
            out[i][tuple(slice1)] = img[tuple(slice2)] - img[tuple(slice3)]

            slice1[axis] = -1
            slice2[axis] = -1
            slice3[axis] = -2

            # 1D equivalent -- out[-1] = (f[-1] - f[-2]) / (x[-1] - x[-2])
            out[i][tuple(slice1)] = img[tuple(slice2)] - img[tuple(slice3)]

            # Reset the slice object in this dimension to ":"
            slice1[axis] = slice(None)
            slice2[axis] = slice(None)
            slice3[axis] = slice(None)
            slice4[axis] = slice(None)

            zero_map = (out[i]==0)
            out[i][zero_map] = eps

        return torch.mean(torch.cat(out, dim=1)**2, dim=(1, 2, 3, 4))

    def training_step(self, batch, batch_idx, optimizer_idx):

        if self.current_epoch==0 and batch_idx==0 and optimizer_idx==0:
            print('Training started...')

        # Get input view, PSF and unmatched ground truth of current batch
        self.last_views, last_PSF, last_slicing = batch['view'], batch['psf'], batch['slicing']
        self.slicing_half = last_slicing[0]

        # Train the generator
        if optimizer_idx == 0:

            # Generate images
            self.last_BSB_sharp_pred, self.last_BSB_blurred_with_PSF = self.forward(self.last_views, last_PSF, 'BSB')

            # Cycle loss
            cycle_loss = self.compute_cycle_loss(self.last_views, self.last_BSB_blurred_with_PSF)

            # WGAN loss
            # Full size
            D_full = self.discriminator_full(self.last_BSB_sharp_pred.detach())
            validity_full = torch.ones_like(D_full).cuda(self.last_BSB_sharp_pred.device.index)
            LSGAN_loss_full = self.compute_LSGAN_loss(D_full, validity_full)

            # Half size
            D_half = self.discriminator_half(self.last_BSB_sharp_pred.detach()[:, :, self.slicing_half[0][0]:self.slicing_half[1][0],
                self.slicing_half[0][1]:self.slicing_half[1][1], self.slicing_half[0][2]:self.slicing_half[1][2]])
            validity_half = torch.ones_like(D_half).cuda(self.last_BSB_sharp_pred.device.index)
            LSGAN_loss_half = self.compute_LSGAN_loss(D_half, validity_half)

            LSGAN_loss = 0.5 * (LSGAN_loss_full + LSGAN_loss_half)

            # Gradient loss
            grad_loss = self.compute_gradient_loss(self.last_BSB_sharp_pred)

            total_loss = self.hparams.cycle_loss_weight * cycle_loss + LSGAN_loss + grad_loss
            tensorboard_logs = {'total_loss': total_loss, 'cycle_loss': cycle_loss, 'LSGAN_loss_G': LSGAN_loss, 'grad_loss': grad_loss}

            return OrderedDict({'loss': total_loss, 'log': tensorboard_logs})

        # Train discriminator with full-size tile
        if optimizer_idx == 1:

            # Sharp prediction
            D_sharp_pred = self.discriminator_full(self.last_BSB_sharp_pred.detach())
            validity_sharp_pred = torch.zeros_like(D_sharp_pred).cuda(self.last_BSB_sharp_pred.device.index)
            LSGAN_loss = self.compute_LSGAN_loss(D_sharp_pred, validity_sharp_pred)

            tensorboard_logs = {'LSGAN_loss_full': LSGAN_loss}

            return OrderedDict({'loss': LSGAN_loss, 'log': tensorboard_logs})

        # Train discriminator with half-size tile
        if optimizer_idx == 2:

            # Sharp prediction
            D_sharp_pred = self.discriminator_half(self.last_BSB_sharp_pred.detach()[:, :, self.slicing_half[0][0]:self.slicing_half[1][0],
                self.slicing_half[0][1]:self.slicing_half[1][1], self.slicing_half[0][2]:self.slicing_half[1][2]])
            validity_sharp_pred = torch.zeros_like(D_sharp_pred).cuda(self.last_BSB_sharp_pred.device.index)
            LSGAN_loss = self.compute_LSGAN_loss(D_sharp_pred, validity_sharp_pred)

            tensorboard_logs = {'LSGAN_loss_half': LSGAN_loss}

            return OrderedDict({'loss': LSGAN_loss, 'log': tensorboard_logs})

    def validation_step(self, batch, batch_idx):
        view, PSF, slicing = batch['view'], batch['psf'], batch['slicing']
        slicing_half = slicing[0]
        slicing_quarter = slicing[1]

        # Generate images
        BSB_sharp_pred, BSB_blurred_with_PSF = self.forward(view, PSF, 'BSB')

        # Cycle loss
        cycle_loss = self.compute_cycle_loss(view, BSB_blurred_with_PSF)

        # WGAN loss
        # Full size
        D_full = self.discriminator_full(BSB_sharp_pred)
        validity_full_1 = torch.ones_like(D_full).cuda(BSB_sharp_pred.device.index)
        validity_full_0 = torch.zeros_like(D_full).cuda(BSB_sharp_pred.device.index)
        LSGAN_loss_full_G = self.compute_LSGAN_loss(D_full, validity_full_1)
        LSGAN_loss_full_D = self.compute_LSGAN_loss(D_full, validity_full_0)

        # Half size
        D_half = self.discriminator_half(BSB_sharp_pred[:, :, slicing_half[0][0]:slicing_half[1][0],
            slicing_half[0][1]:slicing_half[1][1], slicing_half[0][2]:slicing_half[1][2]])
        validity_half_1 = torch.ones_like(D_half).cuda(BSB_sharp_pred.device.index)
        validity_half_0 = torch.zeros_like(D_half).cuda(BSB_sharp_pred.device.index)
        LSGAN_loss_half_G = self.compute_LSGAN_loss(D_half, validity_half_1)
        LSGAN_loss_half_D = self.compute_LSGAN_loss(D_half, validity_half_0)

        LSGAN_loss_G = 0.5 * (LSGAN_loss_full_G + LSGAN_loss_half_G)

        # Gradient loss
        grad_loss = self.compute_gradient_loss(BSB_sharp_pred)

        return {'cycle_loss': cycle_loss, 'LSGAN_loss_G': LSGAN_loss_G, 'LSGAN_loss_full_D': LSGAN_loss_full_D, 'LSGAN_loss_half_D': LSGAN_loss_half_D, 'grad_loss': grad_loss}

    def validation_end(self, outputs):
        avg_cycle_loss = torch.stack([x['cycle_loss'] for x in outputs]).mean()
        avg_LSGAN_loss_G = torch.stack([x['LSGAN_loss_G'] for x in outputs]).mean()
        avg_LSGAN_loss_full_D = torch.stack([x['LSGAN_loss_full_D'] for x in outputs]).mean()
        avg_LSGAN_loss_half_D = torch.stack([x['LSGAN_loss_half_D'] for x in outputs]).mean()
        avg_grad_loss = torch.stack([x['grad_loss'] for x in outputs]).mean()

        tensorboard_logs = {'val_cycle_loss': avg_cycle_loss, 'val_LSGAN_loss_G': avg_LSGAN_loss_G, 'val_LSGAN_loss_full_D': avg_LSGAN_loss_full_D, 'val_LSGAN_loss_half_D': avg_LSGAN_loss_half_D, 'val_grad_loss': avg_grad_loss}

        return {'avg_val_loss': avg_cycle_loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        opt_g = RAdam(self.generator.parameters(), lr=self.hparams.learning_rate)
        opt_d_full = RAdam(self.discriminator_full.parameters(), lr=self.hparams.learning_rate)
        opt_d_half = RAdam(self.discriminator_half.parameters(), lr=self.hparams.learning_rate)
        return [opt_g, opt_d_full, opt_d_half], []

    @pl.data_loader
    def train_dataloader(self):
         if self.hparams.train_dir is None:
            return None
         else:
            dataset = Multiview_Dataset(data_dir=self.hparams.train_dir, num_views=self.hparams.num_views, tile_size=self.hparams.tile_size, shuffle=True, data_augmentation=True)
            return DataLoader(dataset, batch_size=self.hparams.batch_size, shuffle=True)

    @pl.data_loader
    def val_dataloader(self):
        if self.hparams.val_dir is None:
            return None
        else:
            dataset = Multiview_Dataset(data_dir=self.hparams.val_dir, num_views=self.hparams.num_views, tile_size=self.hparams.tile_size, shuffle=False, data_augmentation=False)
            return DataLoader(dataset, batch_size=self.hparams.batch_size)

    def on_epoch_end(self):

        slice_idx = self.last_views.size()[2] // 2

        # Log BSB sharp prediction
        BSB_prediction_grid = torchvision.utils.make_grid(torch.cat([self.last_BSB_sharp_pred[:, :, slice_idx, :, :] for i in range(3)], dim=1))
        self.logger.experiment.add_image('BSB_sharp_pred', BSB_prediction_grid, self.current_epoch)

        # Log the input view
        img_grid = torchvision.utils.make_grid(torch.cat([self.last_views[:, 0, slice_idx, :, :].unsqueeze(1) for i in range(3)], dim=1))
        self.logger.experiment.add_image('Raw_input_view', img_grid, self.current_epoch)

    @staticmethod
    def add_model_specific_args(parent_parser):
        """
        Parameters you define here will be available to your model through self.hparams
        """

        parser = ArgumentParser(parents=[parent_parser])

        # Network params
        parser.add_argument('--in_channels', default=4, type=int)
        parser.add_argument('--out_channels', default=1, type=int)
        parser.add_argument('--feat_channels', default=64, type=int)
        parser.add_argument('--norm_method', default='instance', type=str)
        parser.add_argument('--out_activation', default='sigmoid', type=str)
        parser.add_argument('--cycle_loss_function', default='l1', type=str)
        parser.add_argument('--cycle_loss_weight', default=10.0, type=float)

        # Data
        parser.add_argument('--train_dir', default='/datasets/quad_view/train/', type=str)
        parser.add_argument('--val_dir', default='/datasets/quad_view/val/', type=str)
        parser.add_argument('--data_augmentation', default=True, type=bool)
        parser.add_argument('--num_views', default=4, type=int)

        # Training params (opt)
        parser.add_argument('--batch_size', default=1, type=int)
        parser.add_argument('--learning_rate', default=0.0001, type=float)

        return parser


# def viewPicture():
#     file = pickAFile()
#     pict = makePicture(file)
#     pixel = getPixel(pict, 0, 0)
#     pixels = getPixels(pict)
#     print("width is", getWidth(pict), "height is", getHeight(pict))
#     print("The value of the pixel at 0, 0 is: ", pixel)
#     print("A different way to get the pixel at 0, 0 is: ", pixels[0])
#     explore(pict)

# viewPicture()

# program 16, page 69
def negative(picture):
    for px in getPixels(picture):
        red = getRed(px)
        green = getGreen(px)
        blue = getBlue(px)
        negColor = makeColor(red - red, green - green, blue - blue)
        setColor(px, negColor)
    show(picture)
    repaint(picture)

file = pickAFile()
pict = makePicture(file)
negative(pict)


#!/bin/env python
#
# 基本边缘检测
#

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

def sobel_x(gray_im, plot: plt.Axes = None):
    sx_kernel = np.array([[-1,-2,-1],[0,0,0],[1,2,1]])
    gx_im = cv.filter2D(gray_im, -1, sx_kernel)

    abs_gx_im = np.abs(gx_im)

    if plot is not None:
        plot.set_title('Sobel-X abs gradient')
        plot.set_xticks([])
        plot.set_yticks([])
        plot.imshow(abs_gx_im, cmap='gray')
    return abs_gx_im

def sobel_y(gray_im, plot: plt.Axes = None):
    sy_kernel = np.array([[-1,0,1],[-2,0,2],[-1,0,1]])
    gy_im = cv.filter2D(gray_im, -1, sy_kernel)

    abs_gy_im = np.abs(gy_im)

    if plot is not None:
        plot.set_title('Sobel-Y abs gradient')
        plot.set_xticks([])
        plot.set_yticks([])
        plot.imshow(abs_gy_im, cmap='gray')
    return abs_gy_im

if __name__ == '__main__':
    figure = plt.figure()

    gray_im = cv.imread('CH10_Images/Fig1016(a)(building_original).tif', cv.IMREAD_GRAYSCALE)
    gray_im = gray_im.astype(np.float32)
    # cv.normalize(gray_im, gray_im, 0, 1, cv.NORM_MINMAX)

    oplot = figure.add_subplot(221)
    oplot.set_title('origin')
    oplot.set_xticks([])
    oplot.set_yticks([])
    oplot.imshow(gray_im, cmap='gray')

    # x和y方向上的梯度
    g_sx = sobel_x(gray_im)
    g_sy = sobel_y(gray_im)
    # 梯度图
    gradient_im = g_sx + g_sy
    thresh_g = np.max(gradient_im) * 0.33
    thresh_im = np.zeros_like(gradient_im)
    thresh_im[gradient_im>thresh_g] = 1

    gplot = figure.add_subplot(223)
    gplot.set_title('thresh by 33%')
    gplot.set_xticks([])
    gplot.set_yticks([])
    gplot.imshow(thresh_im, cmap='gray')

    blur_im = cv.medianBlur(gray_im, 5)

    bplot = figure.add_subplot(222)
    bplot.set_title('blurred origin')
    bplot.set_xticks([])
    bplot.set_yticks([])
    bplot.imshow(blur_im, cmap='gray')

    # x和y方向上的梯度
    g_sx_blur = sobel_x(blur_im)
    g_sy_blur = sobel_y(blur_im)
    # 梯度图
    blur_gradient_im = g_sx_blur + g_sy_blur
    thresh_bg = np.max(blur_gradient_im) * 0.33
    blur_thresh_im = np.zeros_like(blur_gradient_im)
    blur_thresh_im[blur_gradient_im>thresh_bg] = 1

    bgplot = figure.add_subplot(224)
    bgplot.set_title('blurred - thresh by 33%')
    bgplot.set_xticks([])
    bgplot.set_yticks([])
    bgplot.imshow(blur_thresh_im, cmap='gray')

    plt.show()