from pathlib import Path
import glob
import os
import re
from pipelines import parsers
from qc_database.utils import relatedness2

class DragenGE:

	def __init__(self,
				 results_dir,
				 sample_names,
				 run_id,
				 sample_expected_files= [],
				run_expected_files = [],
				sample_not_expected_files =[],
				run_not_expected_files= [],
				post_sample_files = [],
				run_complete_marker  = 'post_processing/results/post_processing_finished.txt',
				sample_complete_marker  = 'post_processing/results/post_processing_finished.txt'
				):

		self.results_dir = results_dir
		self.sample_names = sample_names
		self.run_id = run_id
		self.run_complete_marker = run_complete_marker
		self.sample_complete_marker = sample_complete_marker
		self.sample_expected_files = sample_expected_files
		self.sample_not_expected_files = sample_not_expected_files
		self.run_expected_files = run_expected_files
		self.run_not_expected_files = run_not_expected_files
		self.post_sample_files = post_sample_files

	def run_is_complete(self):
		"""
		Look for the file nextflow creates on run end
		"""

		results_path = Path(self.results_dir)

		marker = results_path.glob(self.run_complete_marker)

		if len(list(marker)) == 1:

			return True

		return False

	def run_is_valid(self):
		"""
		Read the file nextflow created on run end (post_processing/results/post_processing_finished.txt)

		open it and see if the success or fail mark is there

		"""

		if self.run_is_complete():

			results_path = Path(self.results_dir)

			marker = results_path.glob(self.run_complete_marker)

			marker = list(marker)[0]

			# last line in file
			last_report = ''

			with open(marker, 'r') as outfile:

				for x in outfile:
					last_report = x.strip()

			if 'success' in last_report:

				return True

		return False

	def get_coverage_metrics(self):

		results_path = Path(self.results_dir)

		run_coverage_metrics_dict = {}

		for sample in self.sample_names:

			sample_coverage_metrics_file = results_path.glob(f'post_processing/results/coverage/*{sample}.depth_summary')

			sample_coverage_metrics_file = list(sample_coverage_metrics_file)[0]

			parsed_coverage_metrics = parsers.parse_custom_coverage_metrics(sample_coverage_metrics_file)

			run_coverage_metrics_dict[sample] = parsed_coverage_metrics

		return run_coverage_metrics_dict

	def get_contamination(self):

		results_path = Path(self.results_dir)

		run_contamination_metrics_dict = {}

		for sample in self.sample_names:

			sample_contamination_metrics_file = results_path.glob(f'post_processing/results/contamination/*{sample}_contamination.selfSM')

			sample_contamination_metrics_file = list(sample_contamination_metrics_file)[0]

			parsed_contamination_metrics = parsers.parse_contamination_metrics(sample_contamination_metrics_file)

			run_contamination_metrics_dict[sample] = parsed_contamination_metrics

		return run_contamination_metrics_dict

	def get_sex_metrics(self):

		results_path = Path(self.results_dir)

		run_sex_metrics_dict = {}

		for sample in self.sample_names:

			sample_sex_metrics_file = results_path.glob(f'post_processing/results/sex/*{sample}_calculated_sex.txt')

			try:
				sample_sex_metrics_file = list(sample_sex_metrics_file)[0]

				parsed_sex_metrics = parsers.parse_dragen_sex_file(sample_sex_metrics_file)

				run_sex_metrics_dict[sample] = parsed_sex_metrics

			except:

				run_sex_metrics_dict[sample] = {'sex': 'Unknown'}

		return run_sex_metrics_dict

	def get_variant_calling_metrics(self):

		results_path = Path(self.results_dir)

		variant_metrics_file = results_path.glob(f'{self.run_id}.vc_metrics.csv')

		variant_metrics_file = list(variant_metrics_file)[0]

		parsed_variant_metrics_file = parsers.parse_dragen_vc_metrics_file(variant_metrics_file)

		return parsed_variant_metrics_file

	def get_alignment_metrics(self):

		results_path = Path(self.results_dir)

		run_alignment_metrics_dict = {}

		for sample in self.sample_names:

			alignment_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.mapping_metrics.csv')

			alignment_metrics_file = list(alignment_metrics_file)[0]

			parsed_alignment_metrics = parsers.parse_dragen_alignment_metrics_file(alignment_metrics_file)

			run_alignment_metrics_dict[sample] = parsed_alignment_metrics

		return run_alignment_metrics_dict

	def get_sensitivity(self):

		results_path = Path(self.results_dir)

		sensitivity_file = results_path.glob(f'post_processing/results/sensitivity/{self.run_id}*_sensitivity.txt')

		sensitivity_file = list(sensitivity_file)

		if len(sensitivity_file) == 0:

			return {'sensitivity': None, 'sensitivity_lower_ci': None, 'sensitivity_higher_ci': None}

		sensitivity_file = sensitivity_file[0]

		parsed_sensitivity_file = parsers.parse_sensitivity_file(sensitivity_file)

		return parsed_sensitivity_file

	def get_variant_count_metrics(self):

		results_path = Path(self.results_dir)

		sample_variant_count_dict = {}

		for sample in self.sample_names:

			vcf_file = results_path.glob(f'post_processing/results/annotated_vcf/{self.run_id}*.vcf.gz')

			vcf_file = list(vcf_file)[0]

			vcf_count_metrics = parsers.get_passing_variant_count(vcf_file, [sample])

			sample_variant_count_dict[sample] = vcf_count_metrics

		return sample_variant_count_dict

	def display_cnv_qc_metrics(self):

		try:
			results_path = Path(self.results_dir)

			cnv_metrics_file = results_path.glob(f'post_processing/results/sv_cnv/qc/{self.run_id}.cnv_qc_report.csv')

			cnv_metrics_file = list(cnv_metrics_file)[0]

			if os.path.isfile(cnv_metrics_file):
				return True
			else:
				return False

		except:
			return False

	def get_postprocessing_cnv_qc_metrics(self):

		results_path = Path(self.results_dir)

		cnv_metrics_file = results_path.glob(f'post_processing/results/sv_cnv/qc/{self.run_id}.cnv_qc_report.csv')

		cnv_metrics_file = list(cnv_metrics_file)[0]

		cnv_metrics_qc_dict = parsers.parse_exome_postprocessing_cnv_qc_metrics(cnv_metrics_file)

		return cnv_metrics_qc_dict

	def get_relatedness_metrics(self, min_relatedness_parents, max_relatedness_unrelated, max_relatedness_between_parents, max_child_parent_relatedness):

		results_path = Path(self.results_dir)

		ped_file = results_path.glob(f'post_processing/results/ped/{self.run_id}.ped')

		ped_file = list(ped_file)[0]

		relatedness_file = results_path.glob(f'post_processing/results/relatedness/{self.run_id}.relatedness2')

		relatedness_file = list(relatedness_file)[0]

		parsed_relatedness, parsed_relatedness_comment = relatedness2.relatedness_test(ped_file, relatedness_file,
																					 min_relatedness_parents,
																					 max_relatedness_unrelated,
																					 max_relatedness_between_parents,
																					 max_child_parent_relatedness)
		return parsed_relatedness, parsed_relatedness_comment

class DragenWGS:

	def __init__(self,
				 results_dir,
				 sample_names,
				 run_id,
				 sample_expected_files= ['*.bam',
									'*.mapping_metrics.csv',
									'*.insert-stats.tab'],
				run_expected_files = ['*.hard-filtered.vcf.gz',
									'*.vc_metrics.csv'],
				sample_not_expected_files =[],
				run_not_expected_files= [],
				post_sample_files = [],
				run_complete_marker  = 'post_processing/results/post_processing_finished.txt',
				sample_complete_marker  = 'post_processing_finished.txt'
				):

		self.results_dir = results_dir
		self.sample_names = sample_names
		self.run_id = run_id
		self.run_complete_marker = run_complete_marker
		self.sample_complete_marker = sample_complete_marker
		self.sample_expected_files = sample_expected_files
		self.sample_not_expected_files = sample_not_expected_files
		self.run_expected_files = run_expected_files
		self.run_not_expected_files = run_not_expected_files
		self.post_sample_files = post_sample_files

	def sample_is_complete(self, sample):
		"""
		Look for presence of file indicating that a sample has completed the pipeline.

		For example the output error log.
		"""

		results_path = Path(self.results_dir)

		new_path = False

		for i in self.run_expected_files:

			if 'post_processing' in i:

				new_path = True
				break

		if new_path:

			results_path = results_path.joinpath('post_processing/results/')

		else:

			results_path = results_path.joinpath('results')

		marker = results_path.glob(self.sample_complete_marker)

		if len(list(marker)) >= 1:

			return True

		return False

	def sample_is_valid(self, sample):
		"""
		Look for files which have to be present for a sample level pipeline to have completed \
		correctly.

		Look for file which if present indicate the pipeline has not finished correctly e.g. intermediate files.
		"""

		results_path = Path(self.results_dir)

		sample_path = results_path.joinpath(sample)

		# check files we want to be there are there
		for file in self.sample_expected_files:

			found_file = sample_path.glob(file)

			if len(list(found_file)) == 0:

				return False

		# check file we do not want to be there are not there
		for file in self.sample_not_expected_files:

			found_file = sample_path.glob(file)

			if len(list(found_file)) > 0:

				return False

		for file in self.post_sample_files:

			split = file.split('{sample}')

			joined = f'{sample}'.join(split)

			found_file = results_path.glob(joined)

			if len(list(found_file)) < 1:

				return False

		return True

	def run_is_complete(self):
		"""
		Look for presence of file indicating that a sample has completed the pipeline.

		For example the output error log.
		"""

		results_path = Path(self.results_dir)

		marker = results_path.glob(self.run_complete_marker)

		if len(list(marker)) == 1:

			return True

		return False

	def run_is_valid(self):
		"""
		Read the file nextflow creates at the end of the run (post_processing/results/post_processing_finished.txt and check if it says fail or success.
		"""

		if self.run_is_complete():

			results_path = Path(self.results_dir)

			marker = results_path.glob(self.run_complete_marker)

			marker = list(marker)[0]

			#last line in file
			last_report = ''

			with open(marker, 'r') as outfile:

				for x in outfile:
					last_report = x.strip()

			if 'success' in last_report:

				return True

		return False

	def run_and_samples_complete(self):
			"""
			Return True if all samples are complete and valid and the run
			level is complete and valid. Otherwise return False.

			"""

			# sample level
			for sample in self.sample_names:

				if self.sample_is_complete(sample) == False:

					return False

			# run level
			if self.run_is_complete() == False:

				return False

			return True

	def run_and_samples_valid(self):

		# sample level
		for sample in self.sample_names:

			if self.sample_is_valid(sample) == False:

				return False

		# run level
		if self.run_is_valid() == False:

			return False

		return True

	def get_variant_calling_metrics(self):

		results_path = Path(self.results_dir)

		variant_metrics_file = results_path.glob(f'{self.run_id}.vc_metrics.csv')

		variant_metrics_file = list(variant_metrics_file)[0]

		parsed_variant_metrics_file = parsers.parse_dragen_vc_metrics_file(variant_metrics_file)

		return parsed_variant_metrics_file

	def get_relatedness_metrics(self, min_relatedness_parents, max_relatedness_unrelated, max_relatedness_between_parents, max_child_parent_relatedness):

		results_path = Path(self.results_dir)

		ped_file = results_path.glob(f'post_processing/results/ped/{self.run_id}.ped')

		ped_file = list(ped_file)[0]

		relatedness_file = results_path.glob(f'post_processing/results/relatedness/{self.run_id}.relatedness2')

		relatedness_file = list(relatedness_file)[0]

		parsed_relatedness, parsed_relatedness_comment = relatedness2.relatedness_test(ped_file, relatedness_file,
																					 min_relatedness_parents,
																					 max_relatedness_unrelated,
																					 max_relatedness_between_parents,
																					 max_child_parent_relatedness)
		return parsed_relatedness, parsed_relatedness_comment

	def get_alignment_metrics(self):

		results_path = Path(self.results_dir)

		run_alignment_metrics_dict = {}

		for sample in self.sample_names:

			alignment_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.mapping_metrics.csv')

			alignment_metrics_file = list(alignment_metrics_file)[0]

			parsed_alignment_metrics = parsers.parse_dragen_alignment_metrics_file(alignment_metrics_file)

			run_alignment_metrics_dict[sample] = parsed_alignment_metrics

		return run_alignment_metrics_dict

	def get_wgs_mapping_metrics(self):

		results_path = Path(self.results_dir)

		run_wgs_coverage_metrics_dict = {}

		for sample in self.sample_names:

			wgs_coverage_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.wgs_coverage_metrics.csv')

			wgs_coverage_metrics_file = list(wgs_coverage_metrics_file)[0]

			parsed_wgs_coverage_metrics = parsers.parse_dragen_wgs_coverage_metrics_file(wgs_coverage_metrics_file)

			run_wgs_coverage_metrics_dict[sample] = parsed_wgs_coverage_metrics

		return run_wgs_coverage_metrics_dict

	def get_exonic_mapping_metrics(self):

		results_path = Path(self.results_dir)

		run_wgs_coverage_metrics_dict = {}

		for sample in self.sample_names:

			wgs_coverage_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.qc-coverage-region-1_coverage_metrics.csv')

			wgs_coverage_metrics_file = list(wgs_coverage_metrics_file)[0]

			parsed_wgs_coverage_metrics = parsers.parse_dragen_wgs_coverage_metrics_file(wgs_coverage_metrics_file)

			run_wgs_coverage_metrics_dict[sample] = parsed_wgs_coverage_metrics

		return run_wgs_coverage_metrics_dict

	def get_ploidy_metrics(self):
		"""
		Need this for version 3.7 Dragen VC to get calculated sex.
		"""

		results_path = Path(self.results_dir)

		run_ploidy_metrics_dict = {}

		for sample in self.sample_names:

			run_ploidy_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.ploidy_estimation_metrics.csv')

			if len(list(run_ploidy_metrics_file)) ==1:

				run_ploidy_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.ploidy_estimation_metrics.csv')

				run_ploidy_metrics_file = list(run_ploidy_metrics_file)[0]

				parsed_run_ploidy_metrics = parsers.parse_ploidy_metrics_file(run_ploidy_metrics_file)

				run_ploidy_metrics_dict[sample] = parsed_run_ploidy_metrics

		return run_ploidy_metrics_dict

	def get_cnv_metrics(self):
		"""
		Get sample level CNV metrics from dragen cnv metrics file
		"""

		results_path = Path(self.results_dir)

		dragen_cnv_metrics_dict = {}

		for sample in self.sample_names:

			cnv_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.cnv_metrics.csv')

			if len(list(cnv_metrics_file)) == 1:

				cnv_metrics_file = results_path.joinpath(sample).glob(f'*{sample}.cnv_metrics.csv')

				cnv_file = list(cnv_metrics_file)[0]

				parsed_cnv_metrics = parsers.parse_dragen_cnv_metrics_file(cnv_file)

				dragen_cnv_metrics_dict[sample] = parsed_cnv_metrics

			else:

				#if not CNV metrics file (i.e. NTC) add empty dictionary for sample
				dragen_cnv_metrics_dict[sample] = {}

		return dragen_cnv_metrics_dict


'''

这是关于多维输入的分类问题，8个输入特征，1个输出

'''

import numpy as np
import torch
import matplotlib.pyplot as plt

xy = np.loadtxt('diabetes.csv.gz', delimiter=',', dtype=np.float32)   # delimiter：以逗号为分解，一般1080，2080显卡都用np.float32
x = torch.from_numpy(xy[:750, :-1])
y = torch.from_numpy((xy[:750, [-1]]))     # 虽然y只取一列，但仍然要保证它与x一致，都是二维矩阵

x_test = torch.from_numpy(xy[750:, :-1])    # 取出最后9个样本作为测试集
y_test = torch.from_numpy(xy[750:, [-1]])

class Multilayer(torch.nn.Module):
    def __init__(self):
        super(Multilayer, self).__init__()
        self.linear1 = torch.nn.Linear(8, 6)
        self.linear2 = torch.nn.Linear(6, 4)
        self.linear3 = torch.nn.Linear(4, 2)
        self.linear4 = torch.nn.Linear(2, 1)
        self.sigmiod = torch.nn.Sigmoid()      # 为了保证最后一层输出是平滑的，所以最后一层始终采用sigmoid()

        # self.activate = torch.nn.Sigmoid()   # 尝试多个激活函数，看看各自效果
        # self.activate = torch.nn.ReLU()      # 根据不同的激活函数，灵活地调整lr
        # self.activate = torch.nn.Tanh()
        self.activate = torch.nn.LogSigmoid()

    def forward(self, x):
        x = self.activate(self.linear1(x))
        x = self.activate(self.linear2(x))
        x = self.activate(self.linear3(x))
        x = self.sigmiod(self.linear4(x))

        return x

model = Multilayer()

critzier = torch.nn.BCELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

loss_draw = []

for epoch in range(10000):     # 增加训练次数模型会好很多
    y_pre = model(x)

    loss = critzier(y_pre, y)
    print('epoch: %d, loss: %f' % (epoch, loss))
    loss_draw.append(loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

y_pred = model(x_test)

y_pred = y_pred.data.numpy()       # 输出时，尽量都转换成numpy()
y_test = y_test.data.numpy()

print('prediction         label ')
for k in range(9):
    print(y_pred[k],'      ',  y_test[k])

plt.plot(loss_draw)    # 绘制loss曲线，不管弹窗会不运行，所以放在最后
plt.show()


import urllib
import urllib2
import csv
from tempfile import SpooledTemporaryFile
import panoply

BASE_URL = 'https://api.oneclickretail.com'
DESTINATION = 'ocr_reports_csv'
IDPATTERN = '{week_asin}'
BATCH_SIZE = 200
FETCH_META = False
DEFAULT_WEEKS_BACK = 1
MAX_SIZE = 100 * (1024 * 1024)  # 100mb

class OcrSource(panoply.DataSource):
    """
    One Click Retail data source.
    API Docs (require login) - https://api.oneclickretail.com/api-docs-login
    Implemented endpoints -
    GET /v5/clients/{client_uuid}/reports/export?format=csv
    """

    def __init__(self, source, options):
        super(OcrSource, self).__init__(source, options)

        if not source.get('destination'):
            source['destination'] = DESTINATION

        if not source.get('idpattern'):
            source['idpattern'] = IDPATTERN

        if not source.get('resources'):
            raise Exception('No resources selected')

        self.resource = None
        self.data = None
        self.resources = source.get('resources')
        self.api_key = source.get('apiKey')
        self.weeks = source.get('weeks', DEFAULT_WEEKS_BACK)
        self.clientUUID = source.get('clientUUID')
        self.processed = 0
        self.total = len(self.resources)

    def read(self, batch_size=None):
        try:
            if not self.resource:
                self.resource = self.resources.pop()
                self.processed += 1
        except IndexError:
            # IndexError will occur when attempting to pop and element
            # off an empty list. This indicates that we have already
            # processed all resources.
            return None

        # Send progress message
        progress_msg = 'Fetching data for %s' % self.resource.get('name')
        self.progress(self.processed, self.total, progress_msg)

        # If data remains from the last resource, use it.
        # Otherwise fetch data from the current resource.
        self.data = self.data or self._fetch_resource()

        batch = self._extract_batch(self.data, batch_size)

        # If no data was returned for this batch we might have reached
        # the end of the source - attempt to read more.
        if not batch:
            # This was the last batch for this resource
            self.resource = None
            # Attempt to fetch the next resouce
            return self.read()

        return batch

    def _fetch_resource(self):
        """
        Assemble the api call, execute it and parse the
        csv response as a list of dicts. Returns a dict generator
        """

        qs = self._build_qs()  # Build the query string
        url = self._build_url(qs)  # Build the full url
        fp = self._api_call(url)  # Fetch the data as a file pointer

        # Parse the list of dicts in to a dict generator
        return csv.DictReader(fp)

    def _build_qs(self):
        params = {
            'format': 'csv',
            'meta': FETCH_META,
            'X-API-KEY': self.api_key,
            'weeks_back': self.weeks
        }

        if self.source.get('filterId'):
            params['filter_id'] = self.source['filterId']

        qs = urllib.urlencode(params)
        return qs

    def _build_url(self, qs):
        base = self.resource.get('value') % self.clientUUID
        url = '%(base)s/%(endpoint)s?%(qs)s' % {
            'base': BASE_URL,
            'endpoint': base,
            'qs': qs
        }
        return url

    def _api_call(self, url):
        """
        Returns a SpooledTemporaryFile - this is a file that is initially
        stored in memory but once its size exceedes max_size it will start
        writing to disk. It is used because there is no way of knowing how
        large of a file the api will return.
        We are expecting a csv file.
        """
        self.log('Request URL', url)
        response = urllib2.urlopen(url)

        # the response MUST be a csv file
        content_type = response.info().get('content-type')
        if 'csv' not in content_type:
            raise Exception('ERROR - Non CSV response.')

        # Decode the returned data and replace any characters that generate
        # encoding errors with the default unicode replacement character.
        data = response.read().decode('utf-8', 'replace')

        self.tmp_file = SpooledTemporaryFile(max_size=MAX_SIZE)
        # Force writing the data encoded as utf-8. If we don't do this,
        # python will attempt to write the data as 'ascii' which does not
        # support special characters

        self.tmp_file.write(data.encode('utf-8'))

        # 'rewind' the file pointer in order to
        # read it back durring `_extract_batch()`
        self.tmp_file.seek(0)

        return self.tmp_file

    def _extract_batch(self, data, batch_size):
        """
        Iterates over BATCH_SIZE of data
        returning a list or results
        """

        batch_size = batch_size or BATCH_SIZE

        batch = []
        try:
            for i in range(batch_size):
                batch.append(data.next())
        except StopIteration:
            pass

        return batch


import sqlite3
from typing import Counter

def get_title(title=None):

    db = sqlite3.connect('netflix.db')
    cur = db.cursor()
    cur.execute(f"""SELECT title, country, release_year, listed_in, description
                    FROM netflix
                    WHERE title like '%{title}%'""")
    s = cur.fetchone()
    db.close()

    return {"title": s[0],
 		"country": s[1],
		"release_year": s[2],
		"genre": s[3],
		"description": s[4]}

def get_year_to_year(year_from, year_to):

    db = sqlite3.connect('netflix.db')
    cur = db.cursor()

    cur.execute(f"""SELECT title, release_year
                    FROM netflix
                    WHERE release_year BETWEEN {year_from} and {year_to}
                    ORDER BY release_year, title
                    LIMIT 100""")

    s = cur.fetchall()
    db.close()

    result = []
    for title in s:
        result.append({"title": title[0],
                        "release_year": title[1]})

    return result

def get_by_rat(rating):

    ratings_list = {"children": "'6'",
                "family": "'G', 'PG', 'PG-13'",
                "adult": "'R', 'NC-17'"}

    db = sqlite3.connect('netflix.db')
    cur = db.cursor()

    cur.execute(f"""SELECT title, rating, description
                    FROM netflix
                    WHERE rating IN ({ratings_list[rating]})
                    LIMIT 100""")

    s = cur.fetchall()
    db.close()

    result = []
    for title in s:
        result.append({"title": title[0],
                        "rating": title[1],
                        "description": title[2]})

    return result

def get_genre(genre):

    db = sqlite3.connect('netflix.db')
    cur = db.cursor()

    cur.execute(f"""SELECT title, description
                    FROM netflix
                    WHERE listed_in LIKE '%{genre}%'
                    LIMIT 10""")
    s = cur.fetchall()
    db.close()

    result_list = []
    for film in s:
        result_list.append({
            "title": film[0],
            "description": film[1]
        })

    return result_list

def pairs(actor1, actor2):

    db = sqlite3.connect('netflix.db')
    cur = db.cursor()

    cur.execute(f"""SELECT 'cast'
        FROM netflix
        WHERE 'cast' LIKE '%{actor1}%'
        AND 'cast' LIKE '%{actor2}%' """)

    s = cur.fetchall()
    db.close()

    result_list = []
    for actor in s:
        result_list.extend(actor[0].split(', '))
    counter = Counter(result_list)
    result_list2 = []
    for actor, count in counter.items():
        if actor not in [actor1, actor2] and count > 2:
            result_list2.append(actor)

    return result_list2

def search_movie(film_type, release_year, genre):
    db = sqlite3.connect('netflix.db')
    cur = db.cursor()

    cur.execute(f"""SELECT title, description
            FROM netflix
            WHERE type = '%{film_type}%'
            AND release_year = '%{release_year}%'
            AND listed_in LIKE '%{genre}%'""")

    s = cur.fetchall()
    db.close()

    result_list = []
    for film in s:
        result_list.append({
            "titile": film[0],
            "description": film[1]
        })
    return result_list

print(search_movie('TV Show', 2000, 'Drama'))


import os, time, random
import multiprocessing as mul

def long_time_task(name):
    print("Run task %s (%s)..." % (name, os.getpid()))
    start = time.time()
    time.sleep(3)
    end = time.time()
    print("task %s runs %0.2f seconds" % (name, (end - start)))

if __name__ == '__main__':
    print("parent process %s" % os.getpid())
    pool = mul.Pool(4)
    for i in range(10):
        pool.apply_async(long_time_task, args=(i, ))
    print("waiting for all subpro")
    pool.close()
    pool.join()
    print("done")


import PyQt5

from PyQt5.QtGui import QKeySequence

from PyQt5.QtCore import Qt
from PyQt5.QtCore import QEventLoop

from PyQt5.QtWidgets import QDialog
from PyQt5.QtWidgets import QDialogButtonBox
from PyQt5.QtWidgets import QApplication
from PyQt5.QtWidgets import QMainWindow
from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QShortcut

from PyQt5.QtWidgets import QVBoxLayout
from PyQt5.QtWidgets import QHBoxLayout

from PyQt5.QtWidgets import QPlainTextEdit
from PyQt5.QtWidgets import QPushButton

from debug_tools import getLogger

# Enable debug messages: ( bitwise )
log = getLogger( 127, __name__ )
log( 1, "Importing " + __name__ )

from debug_tools.utilities import wrap_text

from .utilities import get_screen_center
from .utilities import ignore_exceptions
from .utilities import setTextWithoutCleaningHistory

class StringInputDialog(QMainWindow):
    """
        How can I show a PyQt modal dialog and get data out of its controls once its closed?
        https://stackoverflow.com/questions/18196799/how-can-i-show-a-pyqt-modal-dialog-and-get-data-out-of-its-controls-once-its-clo

        how to clear child window reference stored in parent application when child window is closed?
        https://stackoverflow.com/questions/27420338/how-to-clear-child-window-reference-stored-in-parent-application-when-child-wind
    """

    def __init__(self, parent, settings, fontOptions, _openFileCall, dialogTypeName, dialogTitleMessage):
        super().__init__( parent )
        self.result = 0
        self.setAttribute( Qt.WA_DeleteOnClose )

        self.settings = settings
        self._openFileCall = _openFileCall

        # QWidget::setLayout: Attempting to set QLayout “” on ProgramWindow “”, which already has a layout
        # https://stackoverflow.com/questions/50176661/qwidgetsetlayout-attempting-to-set-qlayout-on-programwindow-which-alre
        self.centralwidget = QWidget()
        self.setCentralWidget( self.centralwidget )
        self.setWindowTitle( dialogTitleMessage )

        # Initial window size/pos last saved. Use default values for first time
        windowScreenGeometry = self.settings.value( "stringInputWindowScreenGeometry" )
        windowScreenState = self.settings.value( "stringInputWindowScreenState" )

        if windowScreenGeometry:
            self.restoreGeometry( windowScreenGeometry )

        else:
            self.resize( 600, 400 )
            # self.move( get_screen_center( self ) )

        if windowScreenState:
            self.restoreState( windowScreenState )

        # nice widget for editing the date
        self.textEditWidget = QPlainTextEdit( self )
        self.textEditWidget.setLineWrapMode( QPlainTextEdit.NoWrap )

        # Detect Ctrl+S ion QTextedit?
        # https://stackoverflow.com/questions/43010630/detect-ctrls-ion-qtextedit
        enterShortcut = QShortcut( QKeySequence( "Ctrl+Enter" ), self.textEditWidget )
        returnShortcut = QShortcut( QKeySequence( "Ctrl+Return" ), self.textEditWidget )
        enterShortcut.activated.connect( self.accept )
        returnShortcut.activated.connect( self.accept )

        # Set initial value of text
        self.textEditWidget.document().setPlainText( wrap_text( """
            # Write your %s here
            # or open it from a new file
        """ % dialogTypeName ) + "\n\n" )

        # Change font, colour of text entry box
        self.textEditWidget.selectAll()
        self.textEditWidget.setStyleSheet( fontOptions )

        # OK and Cancel buttons
        self.standardButtons = QDialogButtonBox( QDialogButtonBox.Ok | QDialogButtonBox.Cancel, Qt.Horizontal, self )
        self.openFileButton = QPushButton( "Open a %s from a file" % dialogTypeName )

        self.standardButtons.accepted.connect( self.accept )
        self.standardButtons.rejected.connect( self.reject )
        self.openFileButton.clicked.connect( self.handleOpenFileCall )

        # Setup the main layout
        self.verticalLayout = QVBoxLayout( self.centralwidget )
        self.horizontalLayout = QHBoxLayout()

        self.verticalLayout.addWidget( self.textEditWidget )
        self.verticalLayout.addLayout( self.horizontalLayout )

        self.horizontalLayout.addWidget( self.openFileButton )
        self.horizontalLayout.addWidget( self.standardButtons )

    def accept(self):
        # log( 1, "Print done" )
        self.result = QDialog.Accepted
        self.close()

    def reject(self):
        # log( 1, "Print reject" )
        self.result = QDialog.Rejected
        self.close()

    def exec_(self):
        """
            Problems with connect in pyqt5
            https://stackoverflow.com/questions/33236358/problems-with-connect-in-pyqt5
        """
        # log( 1, "Blocking main ui" )
        self.event = QEventLoop()
        self.setWindowModality( Qt.ApplicationModal )

        self.show()
        self.event.exec()
        return self.result

    def closeEvent(self, event=None):
        # log( 1, "closeEvent" )
        self.settings.setValue( "stringInputWindowScreenGeometry", self.saveGeometry() )
        self.settings.setValue( "stringInputWindowScreenState", self.saveState() )

        # super().closeEvent( event )
        # self.event.quit()

    def keyPressEvent(self, event):
        """
            Closes the application when the escape key is pressed.
        """

        if event.key() == PyQt5.QtCore.Qt.Key_Escape:
            self.close()

    @ignore_exceptions
    def handleOpenFileCall(self, qt_decorator_bug):
        inputString = self._openFileCall()
        setTextWithoutCleaningHistory( self.textEditWidget, inputString )

    # static method to create the dialog and return ( date, time, accepted )
    @staticmethod
    def getNewUserInput(parent, fontOptions, _openFileCall, dialogTypeName, dialogTitleMessage):
        dialog = StringInputDialog( parent, fontOptions, _openFileCall, dialogTypeName, dialogTitleMessage )
        result = dialog.exec_()

        # dialog.deleteLater()
        # log( 1, "result: %s", result )
        return ( dialog.textEditWidget.toPlainText(), result == QDialog.Accepted )


# -*- coding: utf-8 -*-
# @Date    : 2023-04-27 13:31
# @Author  : chenxuepeng
import sqlite3
from pymysql.converters import escape_string

conn = sqlite3.connect("spider.db", timeout=10, check_same_thread=False)

def test():
    try:
        sql = '''select * from "main"."qt_spider" limit 1'''
        conn.execute(sql)
    except:
        sql = """
        CREATE TABLE "qt_spider" (
      "id" INTEGER NOT NULL,
      "url" TEXT,
      "code" text,
      "req_headers" TEXT,
      "res_headers" TEXT,
      "res_len" TEXT,
      "res_text" TEXT,
      "date" TEXT DEFAULT CURRENT_TIMESTAMP,
      "task" TEXT,
      PRIMARY KEY ("id")
    );"""
        conn.execute(sql)
        conn.commit()

def get_data():
    sql = '''select * from "main"."qt_spider" '''
    fall = conn.execute(sql).fetchall()
    return fall

def del_data():
    sql = '''DELETE from "main"."qt_spider" '''
    fall = conn.execute(sql).fetchall()
    return fall

def get_one_data(ids):
    sql = '''select * from "main"."qt_spider" where id = ?'''
    fall = conn.execute(sql, (ids,)).fetchall()
    return fall

test()


import PyPDF2
import os

def rotate_pdf(path, degrees=90):

    outpath = str(degrees) + '_' + path
    with open(path, 'rb') as in_pdf:

        reader = PyPDF2.PdfFileReader(in_pdf)
        writer = PyPDF2.PdfFileWriter()

        for pnum in range(reader.numPages):
            page = reader.getPage(pnum)
            page.rotateClockwise(degrees)
            writer.addPage(page)

        with open(outpath, 'wb') as out_pdf:
            writer.write(out_pdf)

    print(f"'{path}' rotated {degrees} degrees")

    return True

def get_format(filename):
    name, extension = os.path.splitext(filename)
    return (name, extension[1:].lower())


import frappe
import requests
from datetime import datetime, timedelta
import json

@frappe.whitelist()
def sync_clocify(from_date):
    naive_dt = datetime.strptime(from_date, '%Y-%m-%d')
    naive_dt = naive_dt.replace(hour=0, minute=0, second=1)
    naive_end_dt = naive_dt.replace(hour=23, minute=59, second=59)
    naive_dt = naive_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    naive_end_dt = naive_end_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    # return naive_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    workspace_id =  frappe.db.get_single_value('Clockify Settings', 'workspace_id')
    employee = frappe.db.get_value("Employee", {"user_id": frappe.session.user}, "name")

    if employee:
        user_id = frappe.db.get_value("Clockify Users", {"erp_user": employee}, ["clockify_id", "clockify_key"], as_dict=1)
        if user_id:
            str_url = "https://api.clockify.me/api/v1/workspaces/{workspace_id}/user/{user_id}/time-entries?start={date}&end={end_dt}".format(workspace_id = workspace_id, user_id = user_id.clockify_id, date = naive_dt, end_dt = naive_end_dt)

            url = str_url
            body = {
            }

            header={
                "X-Api-Key": user_id.clockify_key,
                "Content-Type": "application/json",
                "Accept": "application/json"
            }

            time_sheet = []
            # return body
            req = requests.get(url=url, data=body, headers=header)
            if req.status_code == 200:
                if req.json():
                    # return req.json()
                    for item in req.json():
                        start_dt = datetime.strptime(item["timeInterval"]["start"], '%Y-%m-%dT%H:%M:%SZ')
                        start_dt = start_dt + timedelta(hours=5)
                        start_dt = start_dt + timedelta(minutes=30)
                        end_dt = datetime.strptime(item["timeInterval"]["end"], '%Y-%m-%dT%H:%M:%SZ')
                        end_dt = end_dt + timedelta(hours=5)
                        end_dt = end_dt + timedelta(minutes=30)

                        delta = end_dt - start_dt
                        sec = delta.total_seconds()
                        hours = sec / (60 * 60)
                        # ts = frappe.new_doc("Timesheet")
                        # ts.employee = employee
                        # ts.append("time_logs", {
                        #     "activity_type" : "Development",
                        #     "from_time" : start_dt,
                        #     "to_time" : end_dt
                        # })

                        # ts.save(ignore_permissions=True)
                        log = {
                            "from_time": start_dt,
                            "to_time": end_dt,
                            "description": item["description"],
                            "diff": float("%0.4f" % (hours))
                        }
                        time_sheet.append(log)

                else:
                    frappe.msgprint("No Timesheet Found for the Selected Date")
                    return

            return time_sheet
        else:
            frappe.throw("No Clocify Id Assigned to the User")

    else:
        frappe.throw("Employee Not Creeated for you")


from typing import List
from typing import Optional

class TreeNode:

    def __init__(self, val=0, left=None, right=None):
         self.val = val
         self.left = left
         self.right = right

# Print the tree
def pre_order_traversal(node: TreeNode, indent):
    # Root
    if node:
        print(f"{indent}'{node.val}'")
    else:
        return
    indent = indent + "-"
    # Left
    pre_order_traversal(node.left, indent)
    # Right
    pre_order_traversal(node.right, indent)

class Solution:

    x_node = []
    y_node = []

    def pre_order_traversal(self, parent, node: TreeNode, depth, x, y):

        # Root
        if node:
            if node.val == x:
                self.x_node.append(parent)
                self.x_node.append(depth)
            if node.val == y:
                self.y_node.append(parent)
                self.y_node.append(depth)
        else:
            return

        depth += 1
        # Left
        self.pre_order_traversal(node, node.left, depth, x, y)
        # Right
        self.pre_order_traversal(node, node.right, depth, x, y)

    def isCousins(self, root: Optional[TreeNode], x: int, y: int) -> bool:
        self.x_node = []
        self.y_node = []
        self.pre_order_traversal(None, root, 0, x, y)
        return self.x_node[0] != self.y_node[0] and self.x_node[1] == self.y_node[1]

def test():
    solution = Solution()
    # test method
    root = TreeNode(3, None, None)
    node2 = TreeNode(2, None, None)
    root.left = node2
    node4 = TreeNode(4, None, None)
    root.right = node4

    node1 = TreeNode(1, None, None)
    node2.left = node1

    node3 = TreeNode(3, None, None)
    node4.left = node3
    node6 = TreeNode(6, None, None)
    node4.right = node6

    node5 = TreeNode(5, None, None)
    node6.left = node5

    node9 = TreeNode(9, None, None)
    node6.right = node9

    #pre_order_traversal(root, "")

    print(solution.isCousins(root, 1, 9))

    root = TreeNode(1, None, None)
    node2 = TreeNode(2, None, None)
    root.left = node2
    node3 = TreeNode(3, None, None)
    root.right = node3

    node4 = TreeNode(4, None, None)
    node2.right = node4

    node5 = TreeNode(5, None, None)
    node3.right = node5

    # pre_order_traversal(root, "")

    print(solution.isCousins(root, 4, 5))

    root = TreeNode(1, None, None)
    node2 = TreeNode(2, None, None)
    root.left = node2
    node3 = TreeNode(3, None, None)
    root.right = node3

    node4 = TreeNode(4, None, None)
    node2.right = node4

    # pre_order_traversal(root, "")
    print(solution.isCousins(root, 2, 3))

test()


import random
import enum

# Привязываем каждое действие пользователя к числу
class Action(enum.IntEnum):
    камень = 0  # камень
    ножницы = 1  # ножницы
    бумага = 2  # бумага

# Функция не принимает никаких аргументов и возвращает выбор пользователя
def get_user_choice():
    while True:  # проверка "на дурака"
        try:
            user_input = input( '\n''Выбери камень [0], ножницы [1] или бумага [2]: ')
            user_choice = int(user_input)
            action = Action(user_choice)
            return action
        except ValueError:
            range_string = f'[0, {len(Action)} -1]'
            print('Некорректный ввод.')
            continue

# Функция возвращает действие, связанное со случайным числом
def get_computer_choice():
    computer_choice = random.randint(0, len(Action) - 1)  # Значения Action начинаются с 0, а счет len() с 1, вычитаем единицу из len(Action)
    action = Action(computer_choice)
    return action

# Выводим на экран выбор пользователя и компьютера
def print_choices(user_choice, computer_choice):
    print(f'Вы выбрали {user_choice.name}')
    print(f'Компьютер выбрал {computer_choice.name}')

# Определяем победителя
def who_win(user_choice, computer_choice):
    if user_choice == computer_choice:
        print('Ничья!')
    elif user_choice == 0 and computer_choice == 1 or  user_choice == 1 and computer_choice == 2 or user_choice == 2 and computer_choice == 0 :  # камень
        print('Вы победили.')
    else:
        print('Вы проиграли.')

# Играем снова
def play_again():
    while True:
        question = input('\n''Сыграем еще раз? - Да / Нет: ')
        if not question or question.lower() in ('да', 'lf', 'yes'):
            return True
        elif question.lower() in ('нет', 'ytn', 'no'):
            return False
        else:
            print("Пожалуйста, введите 'да' или 'нет'")

# Вызов функций
print('Привет! Ты попал в игру "Камень-ножницы-бумага". Давай начем!')

user_choice = get_user_choice()
computer_choice = get_computer_choice()
print_choices(user_choice, computer_choice)
who_win(user_choice, computer_choice)

while play_again():
    user_choice = get_user_choice()
    computer_choice = get_computer_choice()
    print_choices(user_choice, computer_choice)
    who_win(user_choice, computer_choice)


import os
import math
import time
import random
import shutil
import logging
import warnings
import argparse
import builtins
import numpy as np
import tensorboard_logger as tb_logger

import torch
import torch.nn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.multiprocessing as mp

from utils.resnet import *
from utils.utils_algo import *

from utils.fmnist import load_fmnist
from utils.svhn import load_svhn
from utils.cifar10 import load_cifar10
from utils.cifar100 import load_cifar100
from utils.miniImagenet import load_miniImagenet

from resnet_backbone import resnet_b

parser = argparse.ArgumentParser(description='PyTorch implementation training resnets for various datasets)')
parser.add_argument('--dataset', default='cifar10', type=str,
                    choices=['fmnist', 'SVHN', 'cifar10', 'cifar100', 'miniImagenet'],
                    help='dataset name')

parser.add_argument('--exp-dir', default='./experiment', type=str,
                    help='experiment directory for saving checkpoints and logs')

parser.add_argument('--data-dir', default='./data', type=str)

parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18', choices=['resnet18','resnet34'],
                    help='network architecture (only resnet18 used)')

parser.add_argument('-j', '--workers', default=0, type=int,
                    help='number of data loading workers (default: 0)')

parser.add_argument('--epochs', default=50, type=int,
                    help='number of total epochs to run')

parser.add_argument('--start-epoch', default=0, type=int,
                    help='manual epoch number (useful on restarts)')

parser.add_argument('-b', '--batch-size', default=256, type=int,
                    help='mini-batch size (default: 256), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')

parser.add_argument('--lr', '--learning-rate', default=0.05, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')

parser.add_argument('-lr_decay_epochs', type=str, default='99,199,299',
                    help='where to decay lr, can be a list')

parser.add_argument('-lr_decay_rate', type=float, default=0.1,
                    help='decay rate for learning rate')

parser.add_argument('--cosine', action='store_true', default=False,
                    help='use cosine lr schedule')

parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum of SGD solver')

parser.add_argument('--wd', '--weight_decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')

parser.add_argument('--resume', default='', type=str,
                    help='path to latest checkpoint (default: none)')

parser.add_argument('--world-size', default=-1, type=int,
                    help='number of nodes for distributed training')

parser.add_argument('--rank', default=-1, type=int,
                    help='node rank for distributed training')

parser.add_argument('--dist-url', default='tcp://localhost:10001', type=str,
                    help='url used to set up distributed training')

parser.add_argument('--dist-backend', default='nccl', type=str,
                    help='distributed backend')
parser.add_argument('--rate', default=0.4, type=float,help='-1 for feature, 0.x for random')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')

parser.add_argument('--gpu', default=None, type=int,
                    help='GPU id to use.')

parser.add_argument('--multiprocessing-distributed', action='store_true',
                    help='Use multi-processing distributed training to launch '
                         'N processes per node, which has N GPUs. This is the '
                         'fastest way to use PyTorch for either single node or '
                         'multi node data parallel training')

parser.add_argument('--cuda_VISIBLE_DEVICES', default='0', type=str, \
                        help='which gpu(s) can be used for distributed training')

parser.add_argument('--num-class', default=10, type=int,
                    help='number of class')

parser.add_argument('--low-dim', default=128, type=int,
                    help='embedding dimension')

parser.add_argument('--latent-dim', default=512, type=int,
                    help='latent embedding dimension')

args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda_VISIBLE_DEVICES

torch.set_printoptions(precision=2, sci_mode=False)

logging.basicConfig(format='[%(asctime)s] - %(message)s',
    datefmt='%Y/%m/%d %H:%M:%S',
    level=logging.DEBUG,
    handlers=[
        logging.FileHandler('./result/result_'+ args.dataset + '_bs' + str(args.batch_size) + '_p' \
                            + '_seed' + str(args.seed) + '_dis_' + args.dist_url[-5:] + '.log'),
        logging.StreamHandler()
])
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', best_file_name='model_best.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, best_file_name)

def main():

    iterations = args.lr_decay_epochs.split(',')
    args.lr_decay_epochs = list([])
    for it in iterations:
        args.lr_decay_epochs.append(int(it))

    if args.seed is not None:
        warnings.warn('You have chosen to seed training. This will turn on the CUDNN deterministic setting, which can slow down your training considerably! You may see unexpected behavior when restarting from checkpoints.')
        print()

    if args.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely disable data parallelism.')
        print()

    if args.dist_url == "env://" and args.world_size == -1:
        args.world_size = int(os.environ["WORLD_SIZE"])

    args.distributed = args.world_size > 1 or args.multiprocessing_distributed

    model_path = 'ds_{ds}_lr_{lr}_ep_{ep}_arch_{arch}_sd_{seed}_dis_{dis}'.format(
                                            ds=args.dataset,
                                            lr=args.lr,
                                            ep=args.epochs,
                                            arch=args.arch,
                                            seed=args.seed,
                                            dis=args.dist_url[-5:]
                                            )

    args.exp_dir = os.path.join(args.exp_dir, model_path)
    if not os.path.exists(args.exp_dir):
        os.makedirs(args.exp_dir)

    logging.info(args)
    print()

    ngpus_per_node = torch.cuda.device_count()

    if args.multiprocessing_distributed:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = ngpus_per_node * args.world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
    else:
        # Simply call main_worker function
        print("distibuted training is disabled! ")
        main_worker(args.gpu, ngpus_per_node, args)

def main_worker(gpu, ngpus_per_node, args):

    cudnn.benchmark = True
    model_path = 'ds_{ds}_lr_{lr}_ep_{ep}_arch_{arch}_sd_{seed}_dis_{dis}'.format(
                                            ds=args.dataset,
                                            lr=args.lr,
                                            ep=args.epochs,
                                            arch=args.arch,
                                            seed=args.seed,
                                            dis=args.dist_url[-5:]
                                            )
    args.gpu = gpu

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        np.random.seed(args.seed)
        cudnn.deterministic = True

    if args.gpu is not None:
        print("Use GPU: {} for training\n".format(args.gpu))

    # suppress printing if not master
    if args.multiprocessing_distributed and args.gpu != 0:
        def print_pass(*args):
            pass
        builtins.print = print_pass

    if args.distributed:
        if args.dist_url == "env://" and args.rank == -1:
            args.rank = int(os.environ["RANK"])
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + gpu
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)

    if args.gpu==0:
        print("=> creating model '{}'\n".format(args.arch))
    if args.arch == 'resnet18':
        model = resnet_b.resnet(depth=20, n_outputs=args.num_class)

    if args.arch == 'resnet32':
        model = resnet_b.resnet(depth=32, n_outputs=args.num_class)

    # print(model)

    if args.distributed:
        # For multiprocessing distributed, DistributedDataParallel constructor
        # should always set the single device scope, otherwise,
        # DistributedDataParallel will use all available devices.
        if args.gpu is not None:
            torch.cuda.set_device(args.gpu)
            model.cuda(args.gpu)
            # When using a single GPU per process and per
            # DistributedDataParallel, we need to divide the batch size
            # ourselves based on the total number of GPUs we have
            args.batch_size = int(args.batch_size / ngpus_per_node)
            args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        else:
            model.cuda()
            # DistributedDataParallel will divide and allocate batch_size to all
            # available GPUs if device_ids are not set
            model = torch.nn.parallel.DistributedDataParallel(model)
    elif args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
        # comment out the following line for debugging
        #raise NotImplementedError("Only DistributedDataParallel is supported.")
    else:
        # AllGather implementation (batch shuffle, queue update, etc.) in
        # this code only supports DistributedDataParallel.
        raise NotImplementedError("Only DistributedDataParallel is supported.")

    optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    print("optimizer done!")
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            if args.gpu is None:
                checkpoint = torch.load(args.resume)
            else:
                # Map model to be loaded to specified single gpu.
                loc = 'cuda:{}'.format(args.gpu)
                checkpoint = torch.load(args.resume, map_location=loc)
            args.start_epoch = checkpoint['epoch']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    if args.dataset == 'fmnist':
        train_loader, train_partialY_matrix, train_sampler, test_loader = load_fmnist(args)

    elif args.dataset == 'SVHN':
        train_loader, train_partialY_matrix, train_sampler, test_loader = load_svhn(args)

    elif args.dataset == 'cifar10':
        train_loader, train_partialY_matrix, train_sampler, test_loader = load_cifar10(args)

    elif args.dataset == 'cifar100':
        train_loader, train_partialY_matrix, train_sampler, test_loader = load_cifar100(args)

    elif args.dataset == 'miniImagenet':
        train_loader, train_partialY_matrix, train_sampler, test_loader = load_miniImagenet(args)
        print("miniiamgenet loaded!")
    else:
        raise NotImplementedError("You have chosen an unsupported dataset. Please check and try again.")

    if args.gpu==0:
        logging.info('\nAverage candidate num: {}\n'.format(train_partialY_matrix.sum(1).mean()))

    tempY = train_partialY_matrix.sum(dim=1).unsqueeze(1).repeat(1, train_partialY_matrix.shape[1])
    uniform_confidence = train_partialY_matrix.float() / tempY
    uniform_confidence = uniform_confidence.cuda()

    best_acc = 0
    loss_ce = nn.CrossEntropyLoss()
    for epoch in range(args.start_epoch, args.epochs):

        is_best = False

        acc_train_cls = train(train_loader, model, loss_ce, optimizer, epoch, args)

        acc_test, _ = test(model, test_loader, args, epoch)

        if acc_test.item() > best_acc:
            best_acc = acc_test.item()

        save_checkpoint({
                'epoch': epoch + 1,
                'model': model_path,
                'state_dict': model.state_dict(),
                'optimizer' : optimizer.state_dict()
            }, is_best=is_best, filename='{}/checkpoint.pth.tar'.format(args.exp_dir),
            best_file_name='{}/checkpoint_best.pth.tar'.format(args.exp_dir))

        if args.gpu==0:
            with open(os.path.join(args.exp_dir, 'result.log'), 'a+') as f:
                f.write('Epoch {}: Train_Acc {}, Test_Acc {}, Best_Acc {}. (lr:{})\n'.format(\
                 epoch, acc_train_cls.avg, acc_test.item(), best_acc, optimizer.param_groups[0]['lr']))

            logging.info('Epoch {}: Train_Acc {}, Test_Acc {}, Best_Acc {}. (lr:{})\n'.format(\
                  epoch, acc_train_cls.avg, acc_test.item(), best_acc, optimizer.param_groups[0]['lr']))

def train(train_loader, model,loss_ce,optimizer, epoch, args):
    batch_time = AverageMeter('Time', ':1.2f')
    data_time = AverageMeter('Data', ':1.2f')
    acc_cls = AverageMeter('Acc@Cls', ':2.2f')
    loss_cls_log = AverageMeter('Loss@cls', ':2.2f')

    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, acc_cls, loss_cls_log],
        prefix="Epoch: [{}]".format(epoch)
    )

    model.train()

    end = time.time()

    for i, (images_1, images_2, labels, true_labels, index) in enumerate(train_loader):

        data_time.update(time.time() - end)

        X_1, X_2, Y, index = images_1.cuda(), images_2.cuda(), true_labels.cuda(), index.cuda()

        _,p1 = model(X_1)
        _,p2 = model(X_2)

        p = (p1+p2)/2.0

        loss = loss_ce(Y,p)

        loss_cls_log.update(loss.item())

        acc = accuracy(p, Y)
        acc_cls.update(acc[0].item())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        batch_time.update(time.time() - end)
        end = time.time()

        if args.gpu==0:
            if (i % 50 == 0) or ((i + 1) % len(train_loader) == 0):
                logging.info('Epoch:[{0}][{1}/{2}]\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                    'A_cls {Acc_cls.val:.4f} ({Acc_cls.avg:.4f})\t'
                    'L_cls {Loss_cls.val:.4f} ({Loss_cls.avg:.4f})\t'.format(
                        epoch, i + 1, len(train_loader), batch_time=batch_time,
                        data_time=data_time, Acc_cls=acc_cls,
                        Loss_cls=loss_cls_log
                    )
                )

    return acc_cls, loss_cls_log

def test(model, test_loader, args, epoch):
    with torch.no_grad():

        print('\n=====> Evaluation...\n')
        model.eval()

        top1_acc = AverageMeter("Top1")
        top5_acc = AverageMeter("Top5")

        for batch_idx, (images, labels) in enumerate(test_loader):
            images, labels = images.cuda(), labels.cuda()

            _, outputs = model(images)

            acc1, acc5 = accuracy(outputs, labels, topk=(1, 5))

            top1_acc.update(acc1[0])
            top5_acc.update(acc5[0])

            if args.gpu==0:
                if (batch_idx % 10 == 0) or ((batch_idx + 1) % len(test_loader) == 0):
                    logging.info(
                        'Test:[{0}/{1}]\t'
                        'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                        'Prec@5 {top5.val:.3f} ({top5.avg:.3f})\t'\
                        .format(batch_idx + 1, len(test_loader), top1=top1_acc, top5=top5_acc)
                    )

        # average across all processes
        acc_tensors = torch.Tensor([top1_acc.avg, top5_acc.avg]).cuda(args.gpu)

        dist.all_reduce(acc_tensors)

        acc_tensors /= args.world_size

        if args.gpu==0:
            logging.info('Top1 Accuracy is %.2f%%, Top5 Accuracy is %.2f%%\n'%(acc_tensors[0], acc_tensors[1]))

    return acc_tensors[0], acc_tensors[1]

if __name__ == '__main__':
    main()


'''
Created on Jan 7, 2019

@author: wangxing
'''
import gym
import pybullet_envs
import numpy as np
import math
import argparse
import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from tensorboardX import SummaryWriter

import ptan

ENV_ID = "MinitaurBulletEnv-v0"
GAMMA = 0.99
REWARD_STEPS = 2
BATCH_SIZE = 32
LEARNING_RATE = 5e-5
ENTROPY_BETA = 1e-4
HID_SIZE = 128
TEST_ITERS = 1000

class ACModel(nn.Module):
    def __init__(self, obs_size, act_size):
        super(ACModel, self).__init__()
        self.fc = nn.Sequential(
                nn.Linear(obs_size, HID_SIZE),
                nn.ReLU(),
            )
        self.mu = nn.Sequential(
                nn.Linear(HID_SIZE, act_size),
                nn.Tanh(),
            )
        self.sigma_sq = nn.Sequential(
                nn.Linear(HID_SIZE, act_size),
                nn.Softplus(),
            )
        self.state_val= nn.Sequential(
                nn.Linear(HID_SIZE, 1)
            )

    def forward(self, x):
        fc_out = self.fc(x)
        return self.mu(fc_out), self.sigma_sq(fc_out), self.state_val(fc_out)

class ACAgent:
    def __init__(self, model, device='cpu'):
        self.model = model
        self.device = device

    def initial_state(self):
        return None

    def __call__(self, states, agent_states=None):
        states_np = np.array(states, dtype=np.float32)
        states_v = torch.tensor(states_np).to(self.device)
        mu_v, sigma_sq_v, state_vals_v = self.model(states_v)
        mu = mu_v.data.cpu().numpy()
        sigma = torch.sqrt(sigma_sq_v).data.cpu().numpy()
        actions = np.random.normal(mu, sigma)
        actions = np.clip(actions, -1, 1)
        return actions, agent_states

def test_net(net, env, cnt=10, device='cpu'):
    tot_rew = 0.
    steps = 0
    for _ in range(cnt):
        obs = env.reset()
        while True:
            obs_np = np.array([obs], dtype=np.float32)
            obs_v = torch.tensor(obs_np).to(device)
            mu_v = net(obs_v)[0]
            act = mu_v.squeeze(dim=0).data.cpu().numpy()
            act = np.clip(act, -1, 1)
            sp, r, done, _ = env.step(act)
            tot_rew += r
            steps += 1
            if done:
                break
            obs = sp
    return tot_rew/cnt, steps/cnt

def logprob_gaussian(mu_v, sigma_sq_v, acts_v):
    # Gaussian PDF = 1/sqrt(2*pi*sigma^2) * exp( -(x-mu)^2 / (2*sigma^2) )
    # logprob = -log(sqrt(2*pi*sigma^2))  -(x-mu)^2 / (2 * sigma^2)
    lp1 = -torch.log(torch.sqrt(2 * math.pi * sigma_sq_v))
#     lp2 = -(acts_v - mu_v)^2 / (2 * sigma_sq_v.clamp(min=1e-3))
    lp2 = -(acts_v - mu_v) * (acts_v - mu_v) / (2 * sigma_sq_v.clamp(min=1e-3))
    return lp1 + lp2

def unpack_batch_ac(batch, net, last_val_gamma, device='cpu'):
    states = []
    actions = []
    rewards = []
    not_done_idx = []
    last_states = []
    for i, exp in enumerate(batch):
        states.append(exp.state)
        actions.append(exp.action)
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(i)
            last_states.append(exp.last_state)
    states_np = np.array(states, dtype=np.float32)
    states_v = torch.tensor(states_np).to(device)
    actions_v = torch.FloatTensor(actions).to(device)
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        last_states_np = np.array(last_states, dtype=np.float32)
        last_states_v = torch.tensor(last_states_np).to(device)
        last_vals_v = net(last_states_v)[2]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += last_val_gamma * last_vals_np
    qvals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_v, qvals_v

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cuda', default=False, action='store_true', help='Enable CUDA')

    args = parser.parse_args()
    device = torch.device('cuda' if args.cuda else 'cpu')

    save_path = os.path.join('saves', 'ac-')
    os.makedirs(save_path, exist_ok=True)

    env = gym.make(ENV_ID)
    test_env = gym.make(ENV_ID)

    net = ACModel(env.observation_space.shape[0], env.action_space.shape[0]).to(device)
    print(net)
    agent = ACAgent(net, device=device)

    writer = SummaryWriter(comment='-ac-'+ENV_ID)
    exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, GAMMA, steps_count=REWARD_STEPS)

    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)

    batch = []
    best_rew = None
    with ptan.common.utils.RewardTracker(writer) as tracker:
        for step_idx, exp in enumerate(exp_source):
            rewards_steps = exp_source.pop_rewards_steps()

            if step_idx % TEST_ITERS == 0:
                tic = time.time()
                rewards, steps = test_net(net, test_env, device=device)
                print('Test done in %.2f sec, mean episodic rewards=%.3f, mean episodic steps=%d'\
                      %(time.time()-tic, rewards, steps))
                writer.add_scalar('test_rewards', rewards, step_idx)
                writer.add_scalar('test_steps', steps, step_idx)
                if best_rew is None or best_rew < rewards:
                    if best_rew is not None:
                        print("Best episodic reward updated %.3f ==> %.3f"%(best_rew, rewards))
                        name = 'best_%+.3f_%d.dat' % (rewards, step_idx)
                        fname = os.path.join(save_path, name)
                        torch.save(net.state_dict(), fname)
                best_rew = rewards

            batch.append(exp)
            if len(batch) < BATCH_SIZE:
                continue

            states_v, actions_v, qvals_v = unpack_batch_ac(batch, net, GAMMA**REWARD_STEPS, device=device)
            batch.clear()

            optimizer.zero_grad()
            mu_v, sigma_sq_v, state_vals_v = net(states_v)

            loss_sval = F.mse_loss(input=state_vals_v.squeeze(-1), target=qvals_v)

            adv_v = qvals_v.unsqueeze(dim=-1) - state_vals_v.detach()
            logpi = logprob_gaussian(mu_v, sigma_sq_v, actions_v)
            weighted_logpi = adv_v * logpi
            loss_policy = -weighted_logpi.mean()

            entropy = (-(torch.log(2 * math.pi * sigma_sq_v) + 1) / 2 ).mean()
            loss_entropy = ENTROPY_BETA * entropy

            loss = loss_sval + loss_policy + loss_entropy
            loss.backward()
            optimizer.step()

            writer.add_scalar("advantage", np.mean(adv_v.data.cpu().numpy()), step_idx)
            writer.add_scalar("batch_rewards", np.mean(qvals_v.data.cpu().numpy()), step_idx)
            writer.add_scalar("V(s)", np.mean(state_vals_v.data.cpu().numpy()), step_idx)
            writer.add_scalar("loss_entropy", loss_entropy.item(), step_idx)
            writer.add_scalar("loss_policy", loss_policy.item(), step_idx)
            writer.add_scalar("loss_total", loss.item(), step_idx)


from PyQt5 import QtCore
from PyQt5.QtWidgets import QTableWidget, QAbstractItemView, QHeaderView, QTableWidgetItem
from scripts.Cell import Cell

from scripts.Row import Row

class Table(object):
    def __init__(self) -> None:
        """(class)
        Table object. Creates a table wigdet.Has fixed columns.\n
        """
        super().__init__()
        self.column_names = ["Date", "Name",
                             "Unit", "Quantity",  "Rate", "Total Amount"]

        self.table = QTableWidget()
        self.rows = []

    def create_table(self):
        """(method)
        Returns : self.table
        """
        self.table.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.table.setSortingEnabled(True)
        self.table.setColumnCount(len(self.column_names))
        self.table.setRowCount(0)

        for index, name in enumerate(self.column_names):
            self.table.setHorizontalHeaderItem(index, QTableWidgetItem(name))
            self.table.horizontalHeader().setSectionResizeMode(index, QHeaderView.Stretch)

        return self.table

    def add_row(self, row_item: Row):
        """(method)
        Add rows to table.
        Params:
                row_item : Row object
                row_data : list of text for cells
        """
        # Add row to table.
        self.rows.append(row_item)

        # Increase table row count.
        row = self.table.rowCount()
        self.table.setRowCount(row+1)

        for index, cell in enumerate(row_item.get_item()):
            self.set_table_item(row, cell)

    def edit_row(self, current_row: int, row_data: list):
        row_item: Row = self.rows[current_row]
        row_item.set_item(row_data)
        for index, cell in enumerate(row_item.get_item()):
            self.set_table_item(current_row, cell)

    def insert_row(self, row: int, row_data: Row):
        self.rows.insert(row, row_data)
        self.table.insertRow(row)
        for index, cell in enumerate(row_data.get_item()):
            self.set_table_item(row, cell)

    def delete_row(self, current_row: int):
        self.rows.pop(current_row)
        self.table.removeRow(current_row)

    def set_table_item(self, row: int, cell: Cell):
        item = QTableWidgetItem(cell.get_text())
        self.table.setItem(row, cell.get_column(), item)
        self.table.item(row, cell.get_column()).setTextAlignment(
            QtCore.Qt.AlignmentFlag.AlignCenter)
        return item

    def current_row(self):
        return self.table.currentRow()

    def get_rows(self):
        """(method)
        Returns : Table row object list"""
        return self.rows

    def get_table(self):
        return self.table


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 25 13:33:21 2021

@author: ellenboven
"""

import sys
import argparse
import datetime
import yaml
import time

import numpy as np
import torch
import torch.nn as nn
from ignite.engine import Events
from ignite.metrics import Loss, Accuracy
import os
from os import listdir
from os.path import isfile, join
import math

src_path = 'src_path' #path to src directory
sys.path.insert(0, src_path)

from dataset.linedraw import load_line_draw
from model.init import init_model
from training.setup import set_seed_and_device, setup_training, setup_logging
from training.engine import create_rnn_evaluator, run_training
from correlation import get_popcorr, get_pearson_corr, get_pearson_corr_batch, get_coding_space, get_sparsity, get_weight_info

import logging
logging.getLogger("ignite").setLevel(logging.NOTSET)

import torchvision.transforms as trans

def train_model(
        # Model parameters
        model='LSTM', nlayers=1, nhid=50, dropout=0,
        # Data paramters
        seqlen=50, naddends=2, minval=0, maxval=1, train_noise=0.1, test_noise=0.0,
        training_size=10000, testing_size=1000, val_split=0.2, fixdata=False, input_D=1, npoints=1, targetD=2,
        # Training parameters
        epochs=10, batch_size=50, optim='RMSprop', lr=1e-4, l2_norm=0.0, rate_reg=0.0,
        clip=1.0, early_stopping=False, decay_lr=False, lr_scale=0.1,
        lr_decay_patience=10, keep_hidden=False,
        # Replicability and storage
        save='../../data/sims/linedraw/test', seed=18092, no_cuda=False,
        verbose=False, log_interval=10, bptt=None, nseeds=1, bias=None,
        record_grads=False, spars_int = None, synth_nlayers=0, synth_nhid=0, nfibres=0,
        ablation_epoch=None, synth_ablation_epoch=None,
):
    # Set training seed and get device
    device = set_seed_and_device(seed, no_cuda)

    # Load training data
    train_data, test_data, validation_data = load_line_draw(
        training_size=training_size,
        test_size=testing_size,
        batch_size=batch_size,
        train_val_split=val_split,
        seq_len=seqlen,
        num_addends=naddends,
        minval=minval,
        maxval=maxval,
        train_noise_var=train_noise,
        test_noise_var=test_noise,
        fixdata=fixdata,
        input_D=input_D,
        npoints = npoints,
        targetD = targetD
    )

    #if bptt not set then set as sequence length
    if bptt is None:
        bptt = seqlen

    # Initialise model
    input_size, hidden_size, n_responses, model_type = input_D, nhid, targetD, model

    model = init_model(
        model_type=model_type,
        n_layers=nlayers, hidden_size=nhid,
        input_size=input_size, output_size=n_responses,
        device=device,
        dropout=dropout,
        synth_nlayers=synth_nlayers,
        synth_nhid=synth_nhid,
        predict_last=False
    )

    #bias for RNN intialisation (not modeled in neurips2021)
    if bias is not None:
        with torch.no_grad():
            if model_type in ['DNI_LSTM', 'DNI_TANH', 'cDNI_LSTM', 'cDNI_TANH']:
                model.rnn.rnn.weight_hh_l0[:] = model.rnn.rnn.weight_hh_l0[:] * bias
                model.rnn.rnn.weight_ih_l0[:] = model.rnn.rnn.weight_ih_l0[:] * bias

                model.rnn.rnn.bias_hh_l0[:] = model.rnn.rnn.bias_hh_l0[:]*bias
                model.rnn.rnn.bias_ih_l0[:] = model.rnn.rnn.bias_ih_l0[:]*bias
            else:
                model.rnn.weight_hh_l0[:] = model.rnn.weight_hh_l0[:] * bias
                model.rnn.weight_ih_l0[:] = model.rnn.weight_ih_l0[:] * bias

    #sparsify input 'mossy fibre' synthesiser weights if desired (not modeled in neurips2021)
    if model_type in ['LSTM', 'TANH']:
        nfibres = None
    mask = None
    if nfibres is not None:
        shape = model.rnn.backward_interface.synthesizer.input_trigger.weight.shape
        mask = torch.ones(shape)

        if nfibres != -1:
            for i in range(shape[0]):
                connections = np.random.choice(shape[1], nfibres, replace=False)
                mask[i, connections] = 0
        mask = mask.type(torch.bool)
        with torch.no_grad():
            model.rnn.backward_interface.synthesizer.input_trigger.weight[mask] = 0
            print("Initialised all but {} weights to zero".format(nfibres))

        if nfibres == -1:
            mask = None

    # JOP if bptt not set then set as sequence length
    if bptt is None:
        bptt = seqlen

    # Set up the training regime
    setup = setup_training(
        model, validation_data, optim, ['mse'], lr, l2_norm,
        rate_reg, clip, early_stopping, decay_lr, lr_scale, lr_decay_patience,
        keep_hidden, save, device, True, True, bptt, batch_size, False, record_grads,
        spars_int=spars_int, mask=mask,
    )

    trainer, validator, checkpoint, metrics = setup[:4]
    training_tracer, validation_tracer, timer = setup[4:]

    if verbose:
        setup_logging(
            trainer, validator, metrics,
            len(train_data), log_interval
        )

    # Run training
    test_metrics = run_training(
        model=model,
        train_data=train_data,
        trainer=trainer,
        epochs=epochs,
        test_data=test_data,
        metrics=metrics,
        model_checkpoint=checkpoint,
        device=device
    )

    if record_grads:
        test_metrics, grads = test_metrics

    test_mse = test_metrics['mse']

    print('Training ended: test loss {:5.4f}'.format(
        test_mse))

    print('Saving results....')

    # Save traces
#    training_tracer.save(save)
#    validation_tracer.save(save)

    # Save experiment metadata
    model_params = {
        'model_type': model_type,
        'hidden_size': hidden_size,
        'n_layers': nlayers,
        'input_size': input_size,
        'output_size': n_responses,
        'dropout': dropout,
    }

    learning_params = {
        'optimizer': optim,
        'learning-rate': lr,
        'l2-norm': l2_norm,
        'criterion': 'mse',
        'batch_size': batch_size,
    }

    # Save data parameters in a dictionary for testing
    data_params = {
        'seqlen': seqlen, 'naddends': naddends,
        'minval': minval, 'maxval': maxval,
        'train-size': training_size, 'test-size': testing_size,
        'train-noise': train_noise, 'test-noise': test_noise,
        'val-split': val_split, 'keep-hidden': keep_hidden
    }

    meta = {
        'data-params': data_params,
        'model-params': model_params,
        'learning-params': learning_params,
        'info': {
            #'test-score': test_loss,
            'training-time': timer.value(),
            'timestamp': datetime.datetime.now()
        },
        'seed': seed
    }

    with open(save + '\\meta.yaml', mode='w') as f:
        yaml.dump(meta, f)

    print('Done.')

    # JOP
    train_mse = training_tracer.trace
    print('train mse', train_mse)
    val_mse = [vt[0] for vt in validation_tracer.trace]

    if not record_grads:
        return train_mse, val_mse, train_data, test_data, model, trainer
    else:
        return train_mse, val_mse, train_data, test_data, grads, model, trainer

##############################################################################
# PARSE THE INPUT
##############################################################################

start = time.time()
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Train an LSTM variant on the Line Drawing task')

    # Model parameters
    parser.add_argument('--model', type=str, default='LSTM',
                        help='RNN model to use. One of:'
                        '|TANH|DNI_TANH|LSTM|DNI_LSTM')  # JOP
    parser.add_argument('--nlayers', type=int, default=1,
                        help='number of layers')
    parser.add_argument('--nhid', type=int, default=50,
                        help='number of hidden units per layer')
    parser.add_argument('--dropout', type=float, default=0.0,
                        help='the drop rate for each layer of the network')

    # Data parameters
    parser.add_argument('--seqlen', type=int, default=10,
                        help='sequence length')
    parser.add_argument('--naddends', type=int, default=2,
                        help='the number of addends to be unmasked in a sequence'
                             'must be less than the sequence length')
    parser.add_argument('--minval', type=float, default=0.0,
                        help='minimum value of the addends')
    parser.add_argument('--maxval', type=float, default=1.0,
                        help='maximum value of the addends')
    parser.add_argument('--train-noise', type=float, default=0.0,
                        help='variance of the noise to add to the data')
    parser.add_argument('--test-noise', type=float, default=0.0,
                        help='variance of the noise to add to the data')
    parser.add_argument('--training-size', type=int, default=1000,
                        help='size of the randomly created training set')
    parser.add_argument('--testing-size', type=int, default=10000,
                        help='size of the randomly created test set')
    parser.add_argument('--val-split', type=float, default=0.2,
                        help='proportion of trainig data used for validation')
    parser.add_argument('--fixdata', action='store_true',
                        help='flag to keep the data fixed across each epoch')

    # Training parameters
    parser.add_argument('--epochs', type=int, default=100,
                        help='max number of training epochs')
    parser.add_argument('--batch-size', type=int, default=50, metavar='N',
                        help='size of each batch. keep in mind that if the '
                             'training size modulo this quantity is not zero, then'
                             'the it will be increased to create a full batch.')
    parser.add_argument('--optim', type=str, default='adam',
                        help='gradient descent method, supports on of:'
                             'adam|sparseadam|adamax|rmsprop|sgd|adagrad|adadelta')
    parser.add_argument('--lr', type=float, default=0.001,
                        help='initial learning rate')
    parser.add_argument('--l2-norm', type=float, default=0,
                        help='weight of L2 norm')
    parser.add_argument('--rate-reg', type=float, default=0.0,
                        help='regularization factor for hidden variables')
    parser.add_argument('--clip', type=float, default=1,
                        help='gradient clipping')
    parser.add_argument('--early-stopping', action='store_true',
                        help='use early stopping')
    parser.add_argument('--decay_lr', action='store_true',
                        help='if provided will decay the learning after not'
                             'improving for a certain amount of epochs')
    parser.add_argument('--lr-scale', type=float, default=0.5,
                        help='the factor by which learning rate should be scaled'
                             'after the patience has been exhausted')
    parser.add_argument('--lr_decay_patience', type=int, default=10,
                        help='specifies the number of epochs to wait until'
                             'decay is applied before applying the decay factor')
    parser.add_argument('--keep-hidden', action='store_true',
                        help='keep the hidden state values across an epoch'
                             'of training, detaching them from the computation graph'
                             'after each batch for gradient consistency')

    # Replicability and storage
    parser.add_argument('--save', type=str,
                        default='../../data/sims/linedraw', #replace
                        help='path to save the final model')
    parser.add_argument('--seed', type=int, default=18092,
                        help='random seed')
    parser.add_argument('--nseeds', type=int, default=3,
                        help='number of seeds')

    # CUDA
    parser.add_argument('--no-cuda', action='store_true',
                        help='flag to disable CUDA')

    # Print options
    parser.add_argument('--verbose', action='store_true',
                        help='print the progress of training to std output.')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='report interval')

    # JOP
    parser.add_argument('--bptt', type=int, metavar='N',
                        help='truncation size')
    parser.add_argument('--nfibres', type=int, default=-1, metavar='N',
                        help='Number of non-zero weights from hidden to synthesiser (like mossy fibres)')
    parser.add_argument('--record-grads', action='store_true',
                        help='Record hidden activity and corresponding true/synthetic gradients')
    parser.add_argument('--synth-nlayers', type=int, default=1,
                        help='Number of hidden layers for synthesiser')
    parser.add_argument('--synth-nhid', type=int, default=400, metavar='N',
                        help='Number of hidden units in synthesiser')
    parser.add_argument('--input-D', type=int, default=1,
                    help='input dimension')
    parser.add_argument('--targetD', type=int, default=2,
                    help='target dimension') # for now coded that target dimennsion = input_D
    parser.add_argument('--npoints', type=int, default=7,
                    help='npoints')

    parser.add_argument('--bias', type=float, default=None,
                    help='Bias')
    parser.add_argument('--spars_int', type=int, default=None,
                    help='sparse int')
    parser.add_argument('--ablation_epoch', type=int, default=-1, metavar='N',
                        help='when to ablate the synthesiser..')
    parser.add_argument('--synth_ablation_epoch', type=int, default=-1, metavar='N',
                        help='when to ablate the synthesiser learning (IO)..')

    args = parser.parse_args()

    args.model = 'DNI_LSTM'

    args.record_grads = True
    args.epochs = 1

    # as in paper
    args.nhid = 50
    args.synth_nhid = 400
    args.spars_int = 2
    args.npoints=7
    args.bptt =1

    assert args.record_grads, "You're in the record grads file with record_grads as false. Switch it one!"

    # for 7 points inputs range from (-3, 3)
    args.minval=-math.floor(args.npoints/2)
    args.maxval = math.floor(args.npoints/2)
    if len(np.arange(args.minval, args.maxval+1)) > args.npoints:
        args.minval = args.minval+1

    time_x = time.time()

    print('Variables', args)

                    #################TEST MODELS AND PLOT RESULTS###############
    def create_dir(path, folder):

        path = join(path,folder)

        if not os.path.exists(path):
            try:
                os.mkdir(path)

            except OSError:
                print ("Creation of the directory %s failed" % path)
            else:
                print ("Successfully created the directory %s " % path)

        else:
            print ("Directory %s already exists" % path)
        return path

    if args.record_grads:

        save_path = '/Users/ellenboven/OneDrive - University of Bristol/Documents/PhD/Cerebellearning/Computational/Code/'
        vecname = str(args.model) + "_nseeds-" + str(args.nseeds) + "_nhid-" + str(args.nhid) + "_synthnhid-" + str(args.synth_nhid) + "_Inpsize-" + str(args.input_D) + "_T-" + str(args.bptt) + "_.npy"

        today = datetime.datetime.today()
        experiment = 'LD_lines_grad'
        date =today.strftime('%Y%m%d')

        save_path= create_dir(save_path, experiment)
        root= create_dir(save_path, date)

        nhid = args.nhid
        nbatch = args.batch_size
        biases = [1]
        seeds = [1, 2, 1243, 34521, 135235, 236236, 7, 12, 1115, 987][:args.nseeds]#        biases = [1]
        nfibres_list = [None]
        corr_methods = [get_popcorr, get_pearson_corr, get_coding_space, get_sparsity, get_weight_info]

        score_shape = (len(seeds), len(biases), len(nfibres_list), args.epochs) #add 1 for first batch
        scores = np.zeros(score_shape)
        corr_shape = score_shape[:3] + (args.epochs + 1, len(corr_methods)) # x pop/pearson/etc

        corrs_hidd = np.zeros(corr_shape)
        corrs_truegrads = np.zeros(corr_shape)
        corrs_shidds = np.zeros(corr_shape)
        corrs_sgrad = np.zeros(corr_shape)
        corrs_tgrad = np.zeros(corr_shape)

        ncross_comb = 10 #5 choose 2
        corrs_hidd_vs_shidds = np.zeros(corr_shape)
        corrs_hidd_vs_sgrads = np.zeros(corr_shape)
        corrs_shidds_vs_tgrad = np.zeros(corr_shape[:-1] + ((nhid*2)*args.synth_nhid,))
        corrs_sgrads_vs_tgrad = np.zeros(corr_shape[:-1] + ((nhid*2)**2,))
        corrs_hidd_vs_shidds_all = np.zeros(corr_shape[:-1] + ((nhid*2)*args.synth_nhid,))
        all_scores = np.zeros((len(seeds), 2, args.epochs))

        weights_info = np.zeros(score_shape[:2] + (2, 4,4)) #dni/sparsedni, #weights, #measures(mean, absolute mean, absolute mean without zeros, variance)

        for i, seed in enumerate(seeds):
            args.seed = seed
            for j, bias in enumerate(biases):
                args.bias = bias
                for k, nfibres in enumerate(nfibres_list):
                    if nfibres == -1:
                        args.model = 'LSTM'
                    else:
                        args.model = 'DNI_LSTM'
                    args.nfibres = nfibres

                    #run training
                    train_mse, val_mse, train_data, test_data, grads, model, trainer = train_model(**vars(args))
                    all_scores[i, 0, :] = train_mse
                    all_scores[i, 1, :] = val_mse

                    hidds, truegrads_output, truegrads_cell, tgrads_output, tgrads_cell, sgrads, shidds = grads
                    hidds = torch.stack(hidds).squeeze(1)

                    if nfibres != -1:
                        truegrads_output = torch.stack(truegrads_output).squeeze(1)
                        truegrads_cell = torch.stack(truegrads_cell).squeeze(1)
                        tgrads_output = torch.stack(tgrads_output).squeeze(1)
                        tgrads_cell = torch.stack(tgrads_cell).squeeze(1)
                        sgrads = torch.stack(sgrads).squeeze(1)
                        shidds = torch.stack(shidds).squeeze(1)

                    savedir = "{}_bias-{}_seed-{}".format(args.model, args.bias, args.seed)
                    if nfibres is not None:
                        savedir += "_nfibres-{}".format(nfibres)
                    savedir = os.path.join(root, savedir) + "/"
                    if not os.path.exists(savedir):
                        os.makedirs(savedir)

                    """
                    np.save(savedir + 'hidds.npy', hidds.detach().numpy())
                    np.save(savedir + 'sgrads.npy', sgrads.detach().numpy())
                    np.save(savedir + 'truegrads_output.npy', truegrads_output.detach().numpy())
                    np.save(savedir + 'truegrads_cell.npy', truegrads_cell.detach().numpy())
                    np.save(savedir + 'tgrads_output.npy', tgrads_output.detach().numpy())
                    np.save(savedir + 'tgrads_cell.npy', tgrads_cell.detach().numpy())
                    np.save(savedir + 'shidds.npy', shidds.detach().numpy())

                    np.save(savedir + 'scores.npy', all_scores)

                    np.save(savedir + 'biases.npy', np.array(biases))
                    np.save(savedir + 'seeds.npy', np.array(seeds))
                    """

    print("All finished and results saved in", root)


import json
from json.decoder import JSONDecodeError
from e2j2.exceptions import E2j2Exception

def parse(json_file):
    try:
        with open(json_file) as json_file:
            data = json.load(json_file)
    except IOError:
        # Mark as failed
        raise E2j2Exception('IOError raised while reading file')
    except JSONDecodeError:
        raise E2j2Exception('invalid JSON')

    return data


import scipy.io
import networkx as nx
import matplotlib.pyplot as plt
import random
import numpy as np

# Load the .mat file.
mat2 = scipy.io.loadmat('data/Ciao/rating.mat')
# Print the keys of the mat dictionary to see the variables stored in the .mat file.

print(mat2.keys())
ratings = mat2['rating']
# 打乱数据
np.random.shuffle(ratings)

# 取前1000条数据
num_samples = 10000
ratings = ratings[:num_samples, :]

# 获取所需的列数据
userid = ratings[:, 0]
productid = ratings[:, 1]
categoryid = ratings[:, 2]
rating = ratings[:, 3]
helpfulness = ratings[:, 4]

# 以下是可视化的代码，与之前的代码相同
# 取前10个用户的评分数据
unique_users, user_counts = np.unique(userid, return_counts=True)
top_users = unique_users[:10]

plt.figure(figsize=(12, 8))
for i, user in enumerate(top_users):
    user_ratings = rating[userid == user]
    plt.plot(user_ratings, label='User {}'.format(user))

plt.xlabel('Product Index')
plt.ylabel('Rating')
plt.title('Product Ratings by Top 10 Users')
plt.legend()
plt.show()
# 使用条形图可视化平均评分
unique_products, product_counts = np.unique(productid, return_counts=True)
average_ratings = np.zeros_like(unique_products, dtype=float)
for i, product in enumerate(unique_products):
    product_ratings = rating[productid == product]
    average_ratings[i] = np.mean(product_ratings)

plt.figure(figsize=(10, 6))
plt.bar(unique_products, average_ratings)
plt.xlabel('Product ID')
plt.ylabel('Average Rating')
plt.title('Average Rating per Product')
plt.show()

# 使用散点图可视化评分分布
plt.figure(figsize=(10, 6))
plt.scatter(productid, rating, c=userid)
plt.colorbar(label='User ID')
plt.xlabel('Product ID')
plt.ylabel('Rating')
plt.title('Rating Distribution')
plt.show()

# 使用热力图可视化评分情况
unique_products = np.unique(productid)
unique_categories = np.unique(categoryid)
rating_matrix = np.zeros((len(unique_categories), len(unique_products)), dtype=float)
for i, category in enumerate(unique_categories):
    for j, product in enumerate(unique_products):
        product_ratings = rating[(productid == product) & (categoryid == category)]
        if len(product_ratings) > 0:
            rating_matrix[i, j] = np.mean(product_ratings)

plt.figure(figsize=(12, 8))
plt.imshow(rating_matrix, cmap='hot', aspect='auto')
plt.colorbar(label='Rating')
plt.xticks(np.arange(len(unique_products)), unique_products, rotation='vertical')
plt.yticks(np.arange(len(unique_categories)), unique_categories)
plt.xlabel('Product ID')
plt.ylabel('Category ID')
plt.title('Rating Heatmap')
plt.show()


from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.impute import SimpleImputer
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import OrdinalEncoder
from sklearn.preprocessing import normalize
from pandas.plotting import scatter_matrix
import matplotlib.pyplot as plt
import zipfile
import numpy as np
import pandas as pd
import os

local = '/home/ankit/Downloads'

missing_values = ["n/a", "na", "--"]

train_link = os.path.join(local, 'titanic/train.csv')
test_link = os.path.join(local, 'titanic/test.csv')
gend1er_link = os.path.join(local, 'titanic/gender_submission.csv')

#dataframe import
train = pd.read_csv(train_link, na_values=missing_values)
test = pd.read_csv(test_link)
gender = pd.read_csv(gend1er_link)

def trans(data):
    #drop label first and then run this
    train = data

    train.dropna(subset=['Embarked'],inplace=True)

    train = train.reset_index(drop=True)

    #custom transformation
    class custom_Cabin(BaseEstimator, TransformerMixin):
        def __init__(self,data):
            self.data = data
        def fit(data):
            pass
        def transform(data):
            data.fillna('0', inplace=True) # replace nan with 0
            data_ar = data.values # into array
            data_cat = np.where(data_ar != '0',1,data_ar)
            data_int = data_cat.astype('int')
            cols = data.keys()[0]
            return pd.DataFrame(data_int,columns=[cols])

    cabin_custom = custom_Cabin.transform(train['Cabin'])

    # new feature
    train['relatives'] = train['SibSp'] + train['Parch']

    hot = OneHotEncoder(sparse=False)
    rod = OrdinalEncoder()

    train_cat = hot.fit_transform(train[['Sex']])
    sexcat = list(np.unique(train[['Sex']]))
    sex_df = pd.DataFrame(train_cat, columns=sexcat)
    sex_df = sex_df.add_prefix('sex_')

    train.dropna(subset=['Embarked'],inplace=True)
    embarked_cat = rod.fit_transform(train[['Embarked']])
    embarked_df = pd.DataFrame(embarked_cat, columns=['Embarked'])

    class custom_Cabin(BaseEstimator, TransformerMixin):
        def __init__(self,data):
            self.data = data
        def fit(data):
            pass
        def transform(data):
            data.fillna('0', inplace=True) # replace nan with 0
            data_ar = data.values # into array
            data_cat = np.where(data_ar != '0',1,data_ar)
            return data_cat.astype('int')

    cabin_custom = custom_Cabin.transform(train['Cabin'])

    cabin_custom = pd.DataFrame(cabin_custom, columns=['has_cabin'])
    train.drop(['Embarked'], axis=1, inplace=True)
    train = train.join([embarked_df,sex_df,cabin_custom])

    cont = train[['Age','Fare']]
    cont_cols = list(cont.columns)
    simp = SimpleImputer()
    cont = simp.fit_transform(cont)

    cont_nor = normalize(cont)
    cont_nor = pd.DataFrame(cont_nor, columns = cont_cols)
    cont_nor = cont_nor.add_prefix('norm_')

    end1 = pd.concat([train, cont_nor], axis=1)

    end1.drop(['SibSp', 'Parch', 'Ticket', 'Name','Age','Fare','sex_female','Sex', 'Cabin'],axis=1,inplace=True)
    end1.set_index(['PassengerId'], inplace=True)

    return end1

#calling function and modifying data
train_mod = trans(train)
test_mod = trans(test)

train_label = train_mod['Survived']
train_mod.drop(['Survived'], axis=1, inplace=True)

train_mod.isnull().sum()

#parameter of the model
param_grid = [
  {'C': [1, 10, 100, 1000], 'kernel': ['linear']},
  {'C': [1, 10, 100, 1000], 'gamma': [0.001, 0.0001], 'kernel': ['rbf']},
 ]

# =============================================================================
# #support vector machine
# from sklearn import svm
# lg = svm.SVC()
#
# #Grid search cross validation
# from sklearn.model_selection import GridSearchCV
# grid = GridSearchCV(lg,param_grid, verbose=2, n_jobs=-1,cv=3)
#
# grid.fit(train_mod,np.ravel(train_label))
#
# print('Best estimator parameters: ', grid.best_estimator_)
#
# pred = grid.predict(test_mod)
#
# =============================================================================

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import RandomizedSearchCV

randf = RandomForestClassifier()

para = {'bootstrap': [True, False],
 'max_depth': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, None],
 'max_features': ['auto', 'sqrt'],
 'min_samples_leaf': [1, 2, 4],
 'min_samples_split': [2, 5, 10],
 'n_estimators': [200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000]}

rands =  RandomForestClassifier(randf, para, n_iter = 100, cv = 3, verbose=2, random_state=42, n_jobs = -1)

rands.fit(train_mod,np.ravel(train_label))

pred = rands.predict(test_mod)

from sklearn.metrics import confusion_matrix
print('confusion_matrix: ', confusion_matrix(pred,gender['Survived']))

from sklearn.metrics import accuracy_score
print('accuracy_score: ', accuracy_score(pred, gender['Survived']))