import datetime as dt
from calendar import monthrange
from typing import Optional

from hypothesis.errors import InvalidArgument
from hypothesis.internal.conjecture import utils
from hypothesis.internal.validation import check_type, check_valid_interval
from hypothesis.strategies._internal.core import (
    defines_strategy_with_reusable_values,
    deprecated_posargs,
    just,
    none,
)
from hypothesis.strategies._internal.strategies import SearchStrategy

DATENAMES = ("year", "month", "day")
TIMENAMES = ("hour", "minute", "second", "microsecond")

def is_pytz_timezone(tz):
    if not isinstance(tz, dt.tzinfo):
        return False
    module = type(tz).__module__
    return module == "pytz" or module.startswith("pytz.")

def draw_capped_multipart(data, min_value, max_value):
    assert isinstance(min_value, (dt.date, dt.time, dt.datetime))
    assert type(min_value) == type(max_value)
    assert min_value <= max_value
    result = {}
    cap_low, cap_high = True, True
    duration_names_by_type = {
        dt.date: DATENAMES,
        dt.time: TIMENAMES,
        dt.datetime: DATENAMES + TIMENAMES,
    }
    for name in duration_names_by_type[type(min_value)]:
        low = getattr(min_value if cap_low else dt.datetime.min, name)
        high = getattr(max_value if cap_high else dt.datetime.max, name)
        if name == "day" and not cap_high:
            _, high = monthrange(**result)
        if name == "year":
            val = utils.integer_range(data, low, high, 2000)
        else:
            val = utils.integer_range(data, low, high)
        result[name] = val
        cap_low = cap_low and val == low
        cap_high = cap_high and val == high
    return result

class DatetimeStrategy(SearchStrategy):
    def __init__(self, min_value, max_value, timezones_strat):
        assert isinstance(min_value, dt.datetime)
        assert isinstance(max_value, dt.datetime)
        assert min_value.tzinfo is None
        assert max_value.tzinfo is None
        assert min_value <= max_value
        assert isinstance(timezones_strat, SearchStrategy)
        self.min_value = min_value
        self.max_value = max_value
        self.tz_strat = timezones_strat

    def do_draw(self, data):
        result = draw_capped_multipart(data, self.min_value, self.max_value)
        result = dt.datetime(**result)
        tz = data.draw(self.tz_strat)
        try:
            if is_pytz_timezone(tz):
                # Can't just construct; see http://pytz.sourceforge.net
                return tz.normalize(tz.localize(result))
            return result.replace(tzinfo=tz)
        except (ValueError, OverflowError):
            msg = "Failed to draw a datetime between %r and %r with timezone from %r."
            data.note_event(msg % (self.min_value, self.max_value, self.tz_strat))
            data.mark_invalid()

@defines_strategy_with_reusable_values
@deprecated_posargs
def datetimes(
    min_value: dt.datetime = dt.datetime.min,
    max_value: dt.datetime = dt.datetime.max,
    *,
    timezones: SearchStrategy[Optional[dt.tzinfo]] = none()
) -> SearchStrategy[dt.datetime]:
    """datetimes(min_value=datetime.datetime.min, max_value=datetime.datetime.max, *, timezones=none())

    A strategy for generating datetimes, which may be timezone-aware.

    This strategy works by drawing a naive datetime between ``min_value``
    and ``max_value``, which must both be naive (have no timezone).

    ``timezones`` must be a strategy that generates
    :class:`~python:datetime.tzinfo` objects (or None,
    which is valid for naive datetimes).  A value drawn from this strategy
    will be added to a naive datetime, and the resulting tz-aware datetime
    returned.

    .. note::
        tz-aware datetimes from this strategy may be ambiguous or non-existent
        due to daylight savings, leap seconds, timezone and calendar
        adjustments, etc.  This is intentional, as malformed timestamps are a
        common source of bugs.

    :py:func:`hypothesis.extra.pytz.timezones` requires the :pypi:`pytz`
    package, but provides all timezones in the Olsen database.
    :py:func:`hypothesis.extra.dateutil.timezones` requires the
    :pypi:`python-dateutil` package, and similarly provides all timezones
    there.  If you want to allow naive datetimes, combine strategies
    like ``none() | timezones()``.

    Alternatively, you can create a list of the timezones you wish to allow
    (e.g. from the standard library, :pypi:`dateutil <python-dateutil>`,
    or :pypi:`pytz`) and use :py:func:`sampled_from`.

    Examples from this strategy shrink towards midnight on January 1st 2000,
    local time.
    """
    # Why must bounds be naive?  In principle, we could also write a strategy
    # that took aware bounds, but the API and validation is much harder.
    # If you want to generate datetimes between two particular moments in
    # time I suggest (a) just filtering out-of-bounds values; (b) if bounds
    # are very close, draw a value and subtract its UTC offset, handling
    # overflows and nonexistent times; or (c) do something customised to
    # handle datetimes in e.g. a four-microsecond span which is not
    # representable in UTC.  Handling (d), all of the above, leads to a much
    # more complex API for all users and a useful feature for very few.
    check_type(dt.datetime, min_value, "min_value")
    check_type(dt.datetime, max_value, "max_value")
    if min_value.tzinfo is not None:
        raise InvalidArgument("min_value=%r must not have tzinfo" % (min_value,))
    if max_value.tzinfo is not None:
        raise InvalidArgument("max_value=%r must not have tzinfo" % (max_value,))
    check_valid_interval(min_value, max_value, "min_value", "max_value")
    if not isinstance(timezones, SearchStrategy):
        raise InvalidArgument(
            "timezones=%r must be a SearchStrategy that can provide tzinfo "
            "for datetimes (either None or dt.tzinfo objects)" % (timezones,)
        )
    return DatetimeStrategy(min_value, max_value, timezones)

@defines_strategy_with_reusable_values
@deprecated_posargs
def times(
    min_value: dt.time = dt.time.min,
    max_value: dt.time = dt.time.max,
    *,
    timezones: SearchStrategy[Optional[dt.tzinfo]] = none()
) -> SearchStrategy[dt.time]:
    """times(min_value=datetime.time.min, max_value=datetime.time.max, *, timezones=none())

    A strategy for times between ``min_value`` and ``max_value``.

    The ``timezones`` argument is handled as for :py:func:`datetimes`.

    Examples from this strategy shrink towards midnight, with the timezone
    component shrinking as for the strategy that provided it.
    """
    check_type(dt.time, min_value, "min_value")
    check_type(dt.time, max_value, "max_value")
    if min_value.tzinfo is not None:
        raise InvalidArgument("min_value=%r must not have tzinfo" % min_value)
    if max_value.tzinfo is not None:
        raise InvalidArgument("max_value=%r must not have tzinfo" % max_value)
    check_valid_interval(min_value, max_value, "min_value", "max_value")
    day = dt.date(2000, 1, 1)
    return datetimes(
        min_value=dt.datetime.combine(day, min_value),
        max_value=dt.datetime.combine(day, max_value),
        timezones=timezones,
    ).map(lambda t: t.timetz())

class DateStrategy(SearchStrategy):
    def __init__(self, min_value, max_value):
        assert isinstance(min_value, dt.date)
        assert isinstance(max_value, dt.date)
        assert min_value < max_value
        self.min_value = min_value
        self.max_value = max_value

    def do_draw(self, data):
        return dt.date(**draw_capped_multipart(data, self.min_value, self.max_value))

@defines_strategy_with_reusable_values
def dates(
    min_value: dt.date = dt.date.min, max_value: dt.date = dt.date.max
) -> SearchStrategy[dt.date]:
    """dates(min_value=datetime.date.min, max_value=datetime.date.max)

    A strategy for dates between ``min_value`` and ``max_value``.

    Examples from this strategy shrink towards January 1st 2000.
    """
    check_type(dt.date, min_value, "min_value")
    check_type(dt.date, max_value, "max_value")
    check_valid_interval(min_value, max_value, "min_value", "max_value")
    if min_value == max_value:
        return just(min_value)
    return DateStrategy(min_value, max_value)

class TimedeltaStrategy(SearchStrategy):
    def __init__(self, min_value, max_value):
        assert isinstance(min_value, dt.timedelta)
        assert isinstance(max_value, dt.timedelta)
        assert min_value < max_value
        self.min_value = min_value
        self.max_value = max_value

    def do_draw(self, data):
        result = {}
        low_bound = True
        high_bound = True
        for name in ("days", "seconds", "microseconds"):
            low = getattr(self.min_value if low_bound else dt.timedelta.min, name)
            high = getattr(self.max_value if high_bound else dt.timedelta.max, name)
            val = utils.integer_range(data, low, high, 0)
            result[name] = val
            low_bound = low_bound and val == low
            high_bound = high_bound and val == high
        return dt.timedelta(**result)

@defines_strategy_with_reusable_values
def timedeltas(
    min_value: dt.timedelta = dt.timedelta.min,
    max_value: dt.timedelta = dt.timedelta.max,
) -> SearchStrategy[dt.timedelta]:
    """timedeltas(min_value=datetime.timedelta.min, max_value=datetime.timedelta.max)

    A strategy for timedeltas between ``min_value`` and ``max_value``.

    Examples from this strategy shrink towards zero.
    """
    check_type(dt.timedelta, min_value, "min_value")
    check_type(dt.timedelta, max_value, "max_value")
    check_valid_interval(min_value, max_value, "min_value", "max_value")
    if min_value == max_value:
        return just(min_value)
    return TimedeltaStrategy(min_value=min_value, max_value=max_value)


## importing required modules and libraries
import sys, subprocess, csv
try:
    import fitz
except ImportError:
    print("PyMuPDF is not installed. Installing now...")
    try: subprocess.check_call([sys.executable, "-m", "pip", "install", "PyMuPDF"]); print("PyMuPDF installed successfully."); import fitz
    except subprocess.CalledProcessError: print("Failed to install PyMuPDF.\nPlease use internet to install 'PyMuPDF' package."); sys.exit()
try:
    from tkinter import filedialog
except ImportError:
    print('tkinter is not installed. Installing now...')
    try: subprocess.check_call([sys.executable, "-m", "pip", "install", "tkinter"]); print('tkinker intalled successfully.'); from tkinter import filedialog
    except subprocess.CalledProcessError: print('Failed to install tkinter.\nPlease use internet to install \'Tkinter\' package.'); sys.exit()

filepaths = {'answerkey':None,'transcript':None}

## required functions for task
def GetFiles():
    global filepaths
    print('\nSelect Answer Key pdf file')
    AnswerKey = filedialog.askopenfilename(title='Answer Key',filetypes=[('Answer Key', '*.pdf')])
    print('selected file :',AnswerKey)
    print('Select Transcript pdf file')
    Transcript = filedialog.askopenfilename(title='Transcript',filetypes=[('Transcript', '*.pdf')])
    print('selected file :',Transcript,'\n')
    if FileCheck(AnswerKey,'A') and FileCheck(Transcript,'T'):
        filepaths['answerkey']=AnswerKey
        filepaths['transcript']=Transcript
        return filepaths
    else: print('Please select CORRECT Answer Key and Transcript!'); sys.exit()
# no problem below
def FileCheck(file, K):
    try:
        import fitz
        doc = fitz.open(file)
        text = doc[0].get_text().split('\n')
        print(text)
        if K=='A':
            if 'Indian Institute of Technology, Madras - BS in Data Science and Applications' in text:
                return True
        elif K=='T':
            if 'Name' in text and 'QP Set' in text:
                return True
        return False
    except:
        return False

def TransCSV(file):
    trans = open('./csv files/trans.txt','w',newline='')
    write = csv.writer(trans)
    #reading transcript pdf
    doc = fitz.open(file)
    text = ''.join([page.get_text() for page in doc])
    text = [line.strip() for line in text.split('\n')][:-1]
    # writing information
    write.writerow([text[1]])
    write.writerow([text[3][text[3].index('QP'):]])
    for i in range(11,len(text)):
        if i%2!=0 and i!=len(text)-1 and text[i+1]!='Unanswered':
            if text[i][:4]==text[i+1][:4]:
                write.writerow([text[i],'$'.join(text[i+1].split(','))])
            else:
                write.writerow([text[i],text[i+1]])

def AnswerCSV(file):
    def color(num):
        return 'Green' if num==32512 else 'Red' if num==16711680 else 'Other'
    doc = fitz.open(file)
    answer = open('./csv files/key.txt','w',newline='')
    write = csv.writer(answer)
    #writing question paper id in key
    text=doc[0].get_text().strip().split('\n')
    for x in text:
        if 'IIT M' in x and 'QP' in x:
            write.writerow([x[x.index('QP'):].split()[0]])
            break
    #questions data saving
    def add(Question_id,Question_marks,Question_type,COptions,WOptions):
        if Question_id==None: return
        if Question_type in ['MSQ','MCQ']:
            write.writerow([Question_id,Question_marks,Question_type,'$'.join(COptions),'$'.join(WOptions)])
        elif Question_type in ['SA']:
            write.writerow([Question_id,Question_marks,Question_type,':'.join(COptions[0].split(' to ')),'$'.join(WOptions)])
        Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
    Qcount=0;Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
    for i in range(len(doc)):
        page = doc[i]
        blocks = page.get_text("dict", flags=11)["blocks"]
        for b in blocks:  # iterate through the text blocks
            for l in b["lines"]:  # iterate through the text lines
                for s in l["spans"]:  # iterate through the text spans
                    if s['size']==18 and s['text'][:5]!='Group':
                        if Question_id!=None:
                            add(Question_id,Question_marks,Question_type,COptions,WOptions)
                        Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
                        write.writerow([s['text']])
                    elif ('Question Id' in s['text'] and 'COMPREHENSION' not in s['text']):
                        if Question_id!=None: add(Question_id,Question_marks,Question_type,COptions,WOptions)
                        Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
                        Qcount+=1
                        row = s['text'].split(' ')
                        Question_id = row[7];Question_marks=0;Question_type=row[11];COptions=[];WOptions=[]
                    elif 'Correct Marks' in s['text']:
                        Question_marks = s['text'].split()[3]
                    elif color(s['color']) in ['Green','Red']:
                        if len(s['text'])==len('6406531931004. ') and str(s['text'][:4])=='6406531931004. '[:4]:
                            if color(s['color'])=='Green':
                                COptions.append(s['text'][:-2])
                            elif color(s['color'])=='Red':
                                WOptions.append(s['text'][:-2])
                        elif Question_type=='SA' and color(s['color'])=='Green':
                            COptions.append(s['text'])
                            add(Question_id,Question_marks,Question_type,COptions,WOptions)
                            Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
                    if (i== len(doc)-1 and blocks.index(b)==len(blocks)-1):
                        add(Question_id,Question_marks,Question_type,COptions,WOptions)
                        Question_id=None;Question_marks=None;Question_type=None;COptions=[];WOptions=[]
    # print(f'Total no of questions ():{Qcount}')

def CheckCode():
    key = open('./csv files/key.txt')
    trans = open('./csv files/trans.txt')
    Name = trans.readline() # to remove first line
    Akey = key.readline().strip(); Tkey = trans.readline().strip() #selecting QP set codes
    return (Akey, Tkey)

def Evaluate():
    key = open('./csv files/key.txt')
    trans = open('./csv files/trans.txt')
    Name = trans.readline().strip()
    Akey = key.readline().strip(); Tkey = trans.readline().strip()
    if Akey!=Tkey:
        return 'keys not matching'
    Key = [ques.strip() for ques in key]
    Resp = {}
    for line in trans:
        line = line.strip()
        Resp[line.split(',')[0]]=line.split(',')[1]
    print(f'Hey, {Name}. Your scores in each subject are: ')
    #Grouping courses and evaluating
    Course=None;Answerkey=[];result=[]
    for line in Key:
        if len(line.split(','))==1 or Key.index(line)==len(Key)-1:
            if Course!=None:
                res=(Calculate(Course, Answerkey, Resp))
                if res!=None:
                    result.append(res)
            Course=line;Answerkey=[]
        else:
            Answerkey.append(line.split(','))
    return result

def Calculate(Course, SecQs, Resp):
    if SecQs[0][0] not in Resp:
        return None
    Tmarks=0;Smarks=0
    for ques in SecQs:
        Tmarks+=float(ques[1])
        if ques[0] in Resp:
            if ques[2]=='SA':
                if len(ques[3].split(':'))==1:
                    if ques[3]==Resp[ques[0]]:
                        Smarks+=float(ques[1])
                else:
                    # print(float(ques[3].split(':')[0]),float(Answ[ques[0]]),float(ques[3].split(':')[1]))
                    if float(ques[3].split(':')[0]) <= float(Resp[ques[0]]) <= float(ques[3].split(':')[1]):
                        Smarks+=float(ques[1])
            else:
                count=0;total=len(ques[3].split('$'))
                for ans in Resp[ques[0]].split('$'):
                    if ans in ques[4].split('$'):
                        count=0; break
                    if ans in ques[3].split('$'):
                        count+=1
                Smarks+=(count/total)*float(ques[1])
    marks = (Smarks/Tmarks)*100
    return (Course, marks)

def clear():
    trans = open('./csv files/trans.txt','w',newline='')
    trans.writelines('')
    answ = open('./csv files/key.txt','w',newline='')
    answ.writelines('')


from classes import *
from html_fetch import fetch
from itertools import product
import glob
import os
import pandas as pd
import json
import time
from collections import Counter

def gen_course(course_code: str, session: str):
    data = json.loads(fetch(course_code, session))
    if data["payload"] is None:
        return course_code

    course_data = data["payload"]["pageableCourse"]["courses"][0]

    return_course = Course(course_code)

    temp_section_list = []

    for section in course_data["sections"]:
        if check_enrol_control(section):
            section_class = Section(course_code, section["name"], section["linkedMeetingSections"])
            for meeting in section["meetingTimes"]:
                day = meeting["start"]["day"]
                start = meeting["start"]["millisofday"]//1000//60//60
                end = meeting["end"]["millisofday"]//1000//60//60
                section_class.fill_cal(day, start, end)
            if section["type"] == "Lecture":
                for lec in return_course.lectures:
                    if np.array_equal(lec.times, section_class.times):
                        lec.similar_sections.append(section_class)
                        temp_section_list.append(section_class)
                        continue
                if section_class not in temp_section_list:
                    return_course.lectures.append(section_class)
            elif section["type"] == "Tutorial":
                for tut in return_course.tutorials:
                    if np.array_equal(tut.times, section_class.times):
                        tut.similar_sections.append(section_class)
                        temp_section_list.append(section_class)
                        continue
                if section_class not in temp_section_list:
                    return_course.tutorials.append(section_class)
            elif section["type"] == "Practical":
                for pra in return_course.practicals:
                    if np.array_equal(pra.times, section_class.times):
                        pra.similar_sections.append(section_class)
                        temp_section_list.append(section_class)
                        continue
                if section_class not in temp_section_list:
                    return_course.practicals.append(section_class)

    return return_course

def gen_schedules(section_combos: list[list[Section]], filters: FiltersPackage):
    return_list = []
    conflict_list = []
    for sc in section_combos:
        if check_valid_schedule(sc, filters.dow_selected, filters.dow_hard_enforce):
            return_list.append(gen_schedule(sc, filters))
        else:
            conflict_list.extend(find_conflicts(sc))
    return return_list, conflict_list

def gen_schedule(section_combo, filters: FiltersPackage):
    return_sched = np.zeros((5, 12))
    counter = 1

    for section in section_combo:
        test = section.times * counter
        return_sched = return_sched + test
        counter += 1

    return Schedule(return_sched.transpose(), section_combo, filters)

def find_conflicts(sc: list[Section]):
    base = np.zeros((5, 12))

    return_list = []
    for sec in sc:
        base = base + sec.times
        if not (base <= 1).all():
            return_list.append(sec.course_code)
            base[base != 0] = 1
    return return_list

def check_enrol_control(section):
    for i in section["enrolmentControls"]:
        if i["primaryOrg"]["code"] == "ARTSC" or i["primaryOrg"]["code"] == "*":
            return True
    return False

def check_valid_schedule(sections, dow_selected: list[int], dow_hard_enforce: bool):
    result = None
    if not sections:
        return False
    for section in sections:
        ar = np.array(section.times)
        if result is None:
            result = ar
        else:
            result = result+ar
            if dow_hard_enforce:
                for day in dow_selected:
                    if not (result[day] == 0).all():
                        return False
            if not (result <= 1).all():
                return False

        if section.linked_sections:
            for link in section.linked_sections:
                if not any(x.course_code+x.section_code == section.course_code+link for x in sections):
                    return False

    return True

def create_files(schedules, semester):
    to_del = glob.glob(f"{semester}_*.csv") + glob.glob(f"{semester}_*.xlsx")
    for f in to_del:
        os.remove(f)

    for i in range(min(5, len(schedules))):

        np.savetxt(f"{semester}_schedule_{i}.csv", schedules[i].schedule, delimiter=",", fmt="%s")
    all_files = glob.glob(f"{semester}_*.csv")

    writer = pd.ExcelWriter(f"{semester}_schedules.xlsx", engine="xlsxwriter")

    for f in all_files:
        df = pd.read_csv(f)
        df.to_excel(writer, sheet_name=os.path.basename(f), header=False, index=False)
        writer.sheets[os.path.basename(f)].set_column(0, 7, 30)

    writer.save()

def string_code(sec):
    return " " + sec.section_code if not sec.similar_sections else " " + sec.section_code[:3]

def run_program(fall_cl: list[str], winter_cl: list[str], filters: FiltersPackage):
    t = time.time()
    fall_schedules = program_meat(fall_cl, "f", filters)
    if not isinstance(fall_schedules, list):
        return fall_schedules
    else:
        fall_schedules.sort(key=lambda x: x.score, reverse=True)
    winter_schedules = program_meat(winter_cl, "s", filters)
    if not isinstance(winter_schedules, list):
        return winter_schedules
    else:
        winter_schedules.sort(key=lambda x: x.score, reverse=True)

    fall_scores = dict(Counter(x.score for x in fall_schedules))
    winter_scores = dict(Counter(x.score for x in winter_schedules))

    for i in fall_schedules:
        counter = 1
        i.schedule = i.schedule.astype("object")
        for section in i.section_combo:
            i.schedule[i.schedule == counter] = section.course_code + string_code(section)
            counter += 1
        i.schedule[i.schedule == 0] = ""

    for i in winter_schedules:
        counter = 1
        i.schedule = i.schedule.astype("object")

        for section in i.section_combo:
            i.schedule[i.schedule == counter] = section.course_code + string_code(section)
            counter += 1
        i.schedule[i.schedule == 0] = ""

    return [fall_schedules, winter_schedules]

def program_meat(course_list: list[str], session: str, filters: FiltersPackage):
    courses: list[Course] = []
    sections = []

    for course in course_list:
        gen = gen_course(course, session)
        if not isinstance(gen, Course):
            return CustomException("Existence", gen)
        courses.append(gen)
        if gen.lectures:
            sections.append(gen.lectures)
        if gen.tutorials:
            sections.append(gen.tutorials)
        if gen.practicals:
            sections.append(gen.practicals)

    courses.sort(key=lambda x: (len(x.lectures)+len(x.tutorials)))

    section_combos = [list(t) for t in list(product(*sections))]
    schedules = gen_schedules(section_combos, filters)

    if schedules[0] != []:
        schedules[0].sort(key=lambda x: x.score, reverse=True)
        return schedules[0]
    elif schedules[1] == []:
        return schedules[0]
    else:
        temp = Counter(schedules[1])
        return CustomException("Conflict", max(temp, key=temp.get))

def select_schedules(schedules, fall_best_score, winter_best_score):
    fall_selected = []
    winter_selected = []

    if fall_best_score is not None:
        fall_bests = [sch for sch in schedules[0] if sch.score == fall_best_score]
        current_fall = None
        for sched in fall_bests:
            if current_fall is None:
                fall_selected.append(sched)
                current_fall = sched
                continue
            if np.sum(current_fall.schedule != sched.schedule) >= 4:
                fall_selected.append(sched)
                current_fall = sched

    if winter_best_score is not None:
        winter_bests = [sch for sch in schedules[1] if sch.score == winter_best_score]
        current_winter = None
        for sched in winter_bests:
            if current_winter is None:
                winter_selected.append(sched)
                current_winter = sched
                continue
            if np.sum(current_winter.schedule != sched.schedule) >= 4:
                winter_selected.append(sched)
                current_winter = sched

    return [fall_selected, winter_selected]


# coding:utf-8

import json
import os
import random
import string
from model.file_log import FileLog
from view import route, url_for, View, LoginView, AjaxLoginView, AjaxView

@route('/j/upload', name='j_upload')
class UploadHandler(AjaxLoginView):
    def post(self):
        file1 = self.request.files['file'][0]
        original_fname = file1['filename']
        extension = os.path.splitext(original_fname)[1]
        fname = ''.join(random.choice(string.ascii_lowercase + string.digits) for x in range(24))
        final_filename = fname + extension
        output_file = open("static/uploads/" + final_filename, 'wb')
        output_file.write(file1['body'])
        fl = FileLog.new(final_filename, original_fname, self.current_user())
        self.finish(json.dumps(fl.to_dict()))

@route('/j/upload_file_lst', name='j_upload_file_lst')
class UploadHandler(AjaxLoginView):
    def post(self):
        self.finish(json.dumps(FileLog.get_list()))


""" Tests for regressors module.

Tests in ``py.test`` format.
"""

from os.path import dirname, join as pjoin

import numpy as np
import numpy.linalg as npl
from scipy.interpolate import interp1d

from skimage.filters import threshold_otsu
import nibabel as nib

from fmri_designs.regressors import (events2neural_hr, poly_drift,
                                     deltas_at_rows, spm_hrf_dt,
                                     conds2hrf_cols, f_tests, f_tests_3d,
                                     f_for_outliers, outlier_metrics)
from fmri_designs.tmpdirs import dtemporize
from fmri_designs.spm_funcs import spm_hrf

import pytest
from numpy.testing import assert_array_equal, assert_almost_equal

HERE = dirname(__file__)
CONDS = pjoin(HERE, 'conds')

# Example onset, duration, amplitude triplets
COND_TEST = """\
10    6.0    1
20    4.0    2
24    2.0    0.1"""

# Name of file to write for testing
COND_TEST_FNAME = 'cond_test1.txt'

# Work in temporary directory
@dtemporize
def test_events2neural_hr_simple():
    # test events2neural_hr function
    # Write condition test file
    with open(COND_TEST_FNAME, 'wt') as fobj:
        fobj.write(COND_TEST)
    # Read it back
    times_neural = events2neural_hr(COND_TEST_FNAME, 32, dt=2.)
    assert times_neural.shape == (2, 16)
    times, neural = times_neural
    assert_array_equal(times, np.arange(0, 32, 2))
    # Expected values for tr=2, n_trs=16
    expected = np.zeros(16)
    expected[5:8] = 1
    expected[10:12] = 2
    expected[12] = 0.1
    assert_array_equal(neural, expected)
    times, neural = events2neural_hr(COND_TEST_FNAME, 30, dt=1.)
    assert_array_equal(times, np.arange(30))
    # Expected values for tr=1, n_trs=30
    expected = np.zeros(30)
    expected[10:16] = 1
    expected[20:24] = 2
    expected[24:26] = 0.1
    assert_array_equal(neural, expected)
    # Extend duration, more zeros in neural
    times, neural = events2neural_hr(COND_TEST_FNAME, 40, dt=1.)
    assert_array_equal(times, np.arange(40))
    assert_array_equal(neural, np.concatenate((expected, np.zeros(10))))
    # Drop duration, truncates
    times, neural = events2neural_hr(COND_TEST_FNAME, 20, dt=1.)
    assert_array_equal(times, np.arange(20))
    assert_array_equal(neural, expected[:20])
    # dt of 0.1
    times, neural = events2neural_hr(COND_TEST_FNAME, 30, dt=0.1)
    assert_array_equal(times, np.arange(0, 30, 0.1))
    # Expected values for tr=0.1, n_trs=300
    expected = np.zeros(300)
    expected[100:160] = 1
    expected[200:240] = 2
    expected[240:260] = 0.1
    assert_array_equal(neural, expected)
    # 0.1 is the default
    times, neural = events2neural_hr(COND_TEST_FNAME, 30)
    assert_array_equal(times, np.arange(0, 30, 0.1))
    assert_array_equal(neural, expected)

def test_poly_drift():
    # Test polynomial drift
    times = np.arange(0, 20, 2.15)
    n = len(times)
    assert_array_equal(poly_drift(times, 0), np.ones((n, 1)))
    linear = np.linspace(-1, 1, n)
    linear /= npl.norm(linear)
    assert_almost_equal(poly_drift(times, 1), np.c_[linear, np.ones(n)])
    quadratic = (linear ** 2) - (linear ** 2).mean()
    quadratic /= npl.norm(quadratic)
    assert_almost_equal(poly_drift(times, 2),
                        np.c_[quadratic, linear, np.ones(n)])
    cubic = (linear ** 3) - (linear ** 3).mean()
    cubic /= npl.norm(cubic)
    assert_almost_equal(poly_drift(times, 3),
                        np.c_[cubic, quadratic, linear, np.ones(n)])
    d_10 = poly_drift(times, 10)
    exp_sums = np.concatenate((np.zeros(10), [n]))
    assert_almost_equal(np.sum(d_10, axis=0), exp_sums)
    exp_lengths = np.concatenate((np.ones(10), [np.sqrt(n)]))
    assert_almost_equal(np.sqrt(np.sum(d_10 ** 2, axis=0)), exp_lengths)

def test_deltas_at_rows():
    # Test design for deltas as given rows
    for rows, n in (([3, 7, 15], 20),
                    ([1, 2, 9], 12)):
        d = deltas_at_rows(rows, n)
        for col, row in enumerate(rows):
            assert d[row, col] == 1
            d[row, col] = 0
        assert np.all(d == 0)
    with pytest.raises(IndexError):
        deltas_at_rows([3, 12], 12)

def test_spm_hrf_dt():
    # Test SPM HRF at duration and dt
    assert_almost_equal(spm_hrf(np.arange(0, 30, 0.1)),
                        spm_hrf_dt())
    assert_almost_equal(spm_hrf(np.arange(0, 20, 0.1)),
                        spm_hrf_dt(20))
    assert_almost_equal(spm_hrf(np.arange(10)),
                        spm_hrf_dt(10, 1))

def test_conds2hrf_cols():
    # Test function to return design columns for condition files
    cond_fnames = [pjoin(HERE, 'ds114_sub009_t2r1_cond.txt'),
                   pjoin(HERE, 'new_cond.txt')]
    TR = 2.5
    tr_times = np.arange(0, 400, TR)
    hrf_cols = conds2hrf_cols(cond_fnames, tr_times)
    # Now go the slow way round
    hrf = spm_hrf(np.arange(0, 30, 0.1))
    hr_times, neural = events2neural_hr(cond_fnames[0], 410)
    conv = np.convolve(neural, hrf)[:len(neural)]
    interp0 = interp1d(hr_times, conv, bounds_error=False, fill_value=0)
    hr_times, neural = events2neural_hr(cond_fnames[1], 410)
    conv = np.convolve(neural, hrf)[:len(neural)]
    interp1 = interp1d(hr_times, conv, bounds_error=False, fill_value=0)
    hrf_cols_manual = np.c_[interp0(tr_times), interp1(tr_times)]
    assert_almost_equal(hrf_cols, hrf_cols_manual)

def test_f_tests():
    # Test F test routine against results from R
    # See f_tests.R
    x = np.loadtxt(pjoin(HERE, 'x.txt'))
    y1 = np.loadtxt(pjoin(HERE, 'y1.txt'))
    y2 = np.loadtxt(pjoin(HERE, 'y2.txt'))
    Y = np.c_[y1, y2]
    n = len(x)
    X_f = np.ones((n, 2))
    X_f[:, 1] = x
    F, nu_1, nu_2 = f_tests(Y, X_f, np.ones((n, 1)))
    # Test F test results come from the output of R
    exp_f = [0.01197, 0.5955]
    assert_almost_equal(F, exp_f, 4)
    assert (nu_1, nu_2) == (1, 98)
    # Test 3D version
    vol_shape = (2, 3, 4)
    Y_4d = np.zeros(vol_shape + (n,))
    mask = np.zeros(vol_shape, dtype=bool)
    Y_4d[1, 1, 1] = y1
    Y_4d[1, 1, 2] = y2
    mask[1, 1, 1] = True
    mask[1, 1, 2] = True
    exp_f_3d = np.zeros(vol_shape, dtype=float)
    exp_f_3d[mask] = exp_f
    F, nu_1, nu_2 = f_tests_3d(Y_4d, mask, X_f, np.ones((n, 1)))
    assert_almost_equal(F, exp_f_3d, 4)
    assert (nu_1, nu_2) == (1, 98)

def test_f_for_outliers():
    # More or less smoke test for big-picture routine
    img_fname = pjoin(HERE, 'group00_sub04_run1.nii')
    cond_fnames = [pjoin(CONDS, 'group00_sub04_run1_cond1.txt'),
                   pjoin(CONDS, 'group00_sub04_run1_cond2.txt'),
                   pjoin(CONDS, 'group00_sub04_run1_cond3.txt'),
                   pjoin(CONDS, 'group00_sub04_run1_cond4.txt')]
    tr = 2.0
    f_1, f_2, f_3, msk = f_for_outliers(img_fname, cond_fnames, tr, [2, 3])
    # Long manual process
    data = nib.load(img_fname).get_data()
    n_trs = data.shape[-1]
    tr_times = np.arange(n_trs) * tr
    hrf_cols = conds2hrf_cols(cond_fnames, tr_times)
    drift = poly_drift(tr_times, 3)
    outliers = deltas_at_rows([2, 3], n_trs)
    mean = data.mean(axis=-1)
    mask = mean > threshold_otsu(mean)
    assert_array_equal(mask, msk)
    f_hrf, h_n1, h_n2 = f_tests_3d(data, mask, np.c_[hrf_cols, drift], drift)
    assert_array_equal(f_1[0], f_hrf)
    reduced = np.c_[hrf_cols, drift]
    f_outliers, o_n1, o_n2 = f_tests_3d(data, mask,
                                        np.c_[outliers, reduced], reduced)
    reduced = np.c_[outliers, drift]
    f_both, b_n1, b_n2 = f_tests_3d(data, mask,
                                    np.c_[hrf_cols, reduced], reduced)
    assert_array_equal(f_2[0], f_outliers)
    assert_array_equal(f_3[0], f_both)
    # HRF cols 2 and 3 all zero because of truncated data (first 10 vols)
    assert f_1[1:] == (2, n_trs - 6)
    assert f_2[1:] == (2, n_trs - 8)
    assert f_3[1:] == (2, n_trs - 8)

def test_outlier_metrics():
    a = np.random.normal(10, 2, size=(10, 11, 12))
    b = np.random.normal(10, 2, size=(10, 11, 12))
    c = np.random.normal(10, 2, size=(10, 11, 12))
    msk = np.random.normal(10, 2, size=(10, 11, 12)) > 0.2
    n = np.sum(msk)
    first, second, third, d = outlier_metrics((a, 1, 1), (b, 1, 1),
                                              (c, 1, 1), msk)
    assert first == a[msk].sum() / n
    assert second == b[msk].sum() / n
    assert third == c[msk].sum() / n
    assert d == third - first


from typing import List
from collections import deque

class Solution:
    def uniquePathsIII(self, grid: List[List[int]]) -> int:

        row_n, col_n = len(grid), len(grid[0])
        to_check, results = list(), list()
        starting, ending = list(), tuple()
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1)]
        for x in range(row_n):
            for y in range(col_n):
                if grid[x][y] == 0:
                    to_check.append((x, y))
                elif grid[x][y] == 1:
                    starting = [(x, y)]
                elif grid[x][y] == 2:
                    ending = (x, y)

        walk_length = len(to_check) + 1

        def search(visited=None):
            if visited is None:
                visited = starting
            x, y = visited[-1]

            for dx, dy in directions:
                nx, ny = x+dx, y+dy
                if (nx, ny) in to_check and (nx, ny) not in visited:
                    search(visited + [(nx, ny)])
            if len(visited) == walk_length:
                for dx, dy in directions:
                    nx, ny = x + dx, y+dy
                    if 0 <= nx < row_n and 0 <= ny < col_n and (nx, ny) == ending:
                        results.append(visited)
        search()
        return len(results)

if __name__ == "__main__":
    ex = Solution()

    assert ex.uniquePathsIII([[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 2, -1]]) == 2
    assert ex.uniquePathsIII([[0, 1], [2, 0]]) == 0
    assert ex.uniquePathsIII([[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 2]]) == 4
    assert ex.uniquePathsIII([[1], [2]]) == 1
    assert ex.uniquePathsIII(
        [[1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 2, -1, 0, 0]]) == 2


# -*- coding: utf-8 -*-
"""
Created on 2016/10/13 17:06

@author: qiding
"""

import numpy as np
import statsmodels.api as sm

def main():
    data_len = 100000
    std_err = 0.1

    x_var = gen_x_var(data_len, std=1)
    error = gen_error(data_len, std=std_err)
    y_var = func(x_var, error)

    model = ModelLogNormal(y_var, x_var)
    model.fit()

    data_len_oos = 1000
    x_var_oos = gen_x_var(data_len_oos, std=1)
    error_oos = gen_error(data_len_oos, std=std_err)
    y_var_oos = func(x_var_oos, error_oos)
    y_predict_log, y_predict_with_var_raw_y, y_predict_with_var_err, y_predict_without_var = model.predict(x_var_oos)

    r_sq_without_var = model.report_r_sq(y_var_oos, y_predict_without_var)
    r_sq_with_var_raw = model.report_r_sq(y_var_oos, y_predict_with_var_raw_y)
    r_sq_with_var_err = model.report_r_sq(y_var_oos, y_predict_with_var_err)
    r_sq_log = model.report_r_sq(np.log(y_var_oos), y_predict_log)

    # y_predict_log, y_predict_with_var, y_predict_without_var = model.predict(x_var)
    #
    # r_sq_without_var = model.report_r_sq(y_var, y_predict_without_var)
    # r_sq_with_var = model.report_r_sq(y_var, y_predict_with_var)
    # r_sq_log = model.report_r_sq(np.log(y_var), y_predict_log)

    r_sq_in_sample = model.report_r_sq(y_real=np.log(y_var), y_predict=model.y_var_predict_log_in_sample)

    s = 'without: {}, with_raw: {}, with_err: {}, log: {}, in_sample: {}, in_sample2: {}'.format(
        r_sq_without_var, r_sq_with_var_raw, r_sq_with_var_err, r_sq_log, r_sq_in_sample, model.param_reg.rsquared
    )

    print(s)

class ModelLogNormal:
    def __init__(self, y_var, x_var):
        self.y_var = y_var
        self.x_var = x_var
        self.y_var_log = np.log(y_var)
        self.x_var_log = np.log(x_var)

        self.y_var_predict_log_in_sample = None
        self.error_in_sample = None

        self.model = None
        self.param_reg = None

    def fit(self):
        self.model = sm.OLS(endog=self.y_var_log, exog=sm.add_constant(self.x_var_log))
        self.param_reg = self.model.fit()
        self.y_var_predict_log_in_sample = self.predict_log(self.x_var_log)
        self.error_in_sample = self.y_var_log - self.y_var_predict_log_in_sample

    def predict_log(self, x_var_oos_log):
        exog = sm.add_constant(x_var_oos_log)
        y_predict_log = self.model.predict(params=self.param_reg.params, exog=exog)
        return y_predict_log

    def predict(self, x_var_oos):
        x_var_oos_log = np.log(x_var_oos)
        y_predict_log = self.predict_log(x_var_oos_log)
        y_predict_with_var_raw_y = np.exp(y_predict_log + 1 / 2 * (self.y_var_log.std() ** 2))
        y_predict_with_var_err = np.exp(y_predict_log + 1 / 2 * (self.error_in_sample.std() ** 2))
        y_predict_without_var = np.exp(y_predict_log)
        return y_predict_log, y_predict_with_var_raw_y, y_predict_with_var_err, y_predict_without_var

    @classmethod
    def report_r_sq(cls, y_real, y_predict):
        err = y_real - y_predict
        msr = (err * err).mean()
        err2 = y_real - y_real.mean()
        mse = (err2 * err2).mean()
        r_sq = 1 - msr / mse
        return r_sq

def gen_x_var(data_len, std):
    normal_data = np.random.normal(size=data_len, scale=std)
    exp_normal_data = np.exp(normal_data)
    return exp_normal_data

def gen_error(data_len, std):
    err = np.random.normal(scale=std, size=data_len)
    return err

def func(x_var, error, a=.8, b=.4):
    log_x = np.log(x_var)
    log_y = a * log_x + b + error
    y_exp = np.exp(log_y)
    return y_exp

if __name__ == '__main__':
    main()


#!/usr/bin/env python
from setuptools import find_packages, setup

install_requires = [
    'wagtail>=2.0',
]

docs_require = [
    'sphinx>=1.4.0',
]

tests_require = [
    'pytest-cov==2.4.0',
    'pytest-django==3.1.2',
    'pytest==3.0.5',
    'requests-mock==1.1.0',

    # Linting
    'isort==4.2.5',
    'flake8==3.0.3',
    'flake8-blind-except==0.1.1',
    'flake8-debugger==1.4.0',
    'flake8-imports==0.1.0',
]

setup(
    name='wagtail-audit-trail',
    version='1.3.0',
    description="Wagtail audit trail",
    long_description=open('README.rst', 'r').read(),
    url='https://github.com/labd/wagtail-audit-trail',
    author="Lab Digital",
    author_email="opensource@labdigital.nl",
    install_requires=install_requires,
    tests_require=tests_require,
    extras_require={
        'docs': docs_require,
        'test': tests_require,
    },
    entry_points={},
    package_dir={'': 'src'},
    packages=find_packages('src'),
    include_package_data=True,
    license='MIT',
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Environment :: Web Environment',
        'Framework :: Django',
        'Framework :: Django :: 1.11',
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.6',
    ],
    zip_safe=False,
)


import functools
import logging
import unicodedata

class Normaliser:
	"""
	Normalises strings and keeps track of those that (1) do not comply to
	Unicode's normal form; (2) have whitespace issues.
	"""

	def __init__(self, nfc_chars=[]):
		"""
		Constructor. The optional arg specifies the set of chars that should
		not be decomposed.
		"""
		self.log = logging.getLogger(__name__)

		self.norm_f = functools.partial(unicodedata.normalize, 'NFD')

		self.nfc_chars = set(nfc_chars)

		self.strip_errors = []
		self.norm_errors = []

	def normalise(self, string, line_num):
		"""
		Strips the whitespace and applies Unicode normalisation to the given
		string. The second arg is used as an ID of the string when reporting
		its lint errors (if such).
		"""
		stripped = string.strip()
		if stripped != string:
			self.strip_errors.append(line_num)

		nfc_pos = [index
					for index, char in enumerate(stripped)
					if char in self.nfc_chars]

		parts = []
		start_pos = 0

		for pos in nfc_pos:
			if pos > 0:
				parts.append(self.norm_f(stripped[start_pos:pos]))

			parts.append(stripped[pos])
			start_pos = pos + 1

		if start_pos < len(stripped):
			parts.append(self.norm_f(stripped[start_pos:]))

		norm = ''.join(parts)

		if norm != stripped:
			self.norm_errors.append(line_num)

		return norm

	def report(self, reporter, ignore_nfd=False, ignore_ws=False):
		"""
		Adds the problems that have been found so far to the given Reporter
		instance. The two keyword args can be used to restrict the error types
		to be reported.
		"""
		if self.strip_errors and not ignore_ws:
			reporter.add(self.strip_errors, 'leading or trailing whitespace')

		if self.norm_errors and not ignore_nfd:
			reporter.add(self.norm_errors, 'not in Unicode NFD')


import numpy as np
from numba import njit
# import g3read_units as g3u

###----- UNIT HANDLER -----###
class Uregexpr(object):
    def __init__(self, expression):
        self.expression = expression
    def evaluate(self, ureg):
        return ureg.parse_expression(self.expression)
    def __str__(self):
        return 'ureg '+self.expression

# type of input paramters, Uregexpr means the input may be  a string like '1000. kpc'
types={
    "CENTER_X": Uregexpr,
    "CENTER_Y": Uregexpr,
    "CENTER_Z": Uregexpr,
    "IMG_XY_SIZE": Uregexpr,
    "IMG_Z_SIZE": Uregexpr,
    "IMG_SIZE": int,

    }

xyz_keys = ['CENTER_X', 'CENTER_Y', 'CENTER_Z']

def parse_ureg_expressions(kv, to_parse, ureg):
    """
    Reads input values with units (e.g. '3000 kpc') and parses it in a  `pint` value).
    Input values came from the config file and are passed in a dict `kv`.
    `to_parse` is a dictionary that stores the type of the value of a given key (e.g. the number of pixels IMG_SIZE is `int`),
    and `ureg` is a `pint` instance
    """
    res = {}
    for k in to_parse:
        res[k] = kv[k].evaluate(ureg)
    return res

###----- KERNELS -----###
@njit
def WendlandC6_2D(u: float, h_inv: float):
    """
    Evaluate the WendlandC6 spline at position u, where u:= x*h_inv, for a 2D projection.
    Values correspond to 295 neighbours.
    """

    norm2D = 78. / (7. * 3.1415926)

    if u < 1.0:
        n = norm2D * h_inv**2
        u_m1 = (1.0 - u)
        u_m1 = u_m1 * u_m1  # (1.0 - u)^2
        u_m1 = u_m1 * u_m1  # (1.0 - u)^4
        u_m1 = u_m1 * u_m1  # (1.0 - u)^8
        u2 = u*u
        return ( u_m1 * ( 1.0 + 8*u + 25*u2 + 32*u2*u )) * n
    else:
        return 0.

###----- GRID -----###
@njit
def calc_area_weights(dw: np.ndarray,
                masked_x: np.ndarray, masked_y: np.ndarray,
                masked_h: np.ndarray, kernel,
                bin_min: int, delta_bin: float, n_bins: int):

    N = len(masked_x)

    tot_max_bin_spread = 0

    for k in range(N):
        xpx = (masked_x[k]-bin_min)/delta_bin #pos of x in terms of px (from 0 to n_bins-1)
        ypx = (masked_y[k]-bin_min)/delta_bin #pos of y in terms of px
        bin_i = int(xpx) #calculates distance in i-direction from minimum bin
        bin_j = int(ypx) #calculates distance in j-direction
        Hsml = masked_h[k]/delta_bin # in number of px

        max_bin_spread = int(Hsml)

        tot_max_bin_spread+=max_bin_spread
        distr_weight = 0.

        for i in range(bin_i - max_bin_spread, bin_i + max_bin_spread + 1):
            if i>=n_bins or i<0:
                continue
            for j in range(bin_j - max_bin_spread, bin_j + max_bin_spread + 1):
                if j>=n_bins or j<0:
                    continue
                u = ( (i+.5-xpx)**2 + (j+.5-ypx)**2 )**.5 / Hsml # distance of particle to px center
                if u > 1.0:
                    continue

                dx = min(xpx + Hsml, i+1) - max(xpx-Hsml, i)
                dy = min(ypx + Hsml, j+1) - max(ypx-Hsml, j)
                dxdy = dx*dy # area of px

                if Hsml < 1.:
                    wk = dxdy
                else:
                    wk = kernel(u, 1./Hsml) * dxdy

                distr_weight += wk

        dw[k] += distr_weight
    if N>0:
        return tot_max_bin_spread/N
    else:
        return np.nan

@njit
def add_to_grid(final_image_t: np.ndarray, weight_image: np.ndarray,
                masked_x: np.ndarray, masked_y: np.ndarray,
                masked_h: np.ndarray, masked_w: np.ndarray, masked_q: np.ndarray,
                dz: float, kernel, distr_weight: np.ndarray,
                bin_min: int, delta_bin: float, n_bins: int):
    """
    This routine adds a chunk of particles with sky positions `masked_x, masked_y`,
    smoothing length  `masked_h,` and value `masked_w` (e.g. paticle mass) to a FIT buffer `finalt_image_t`.
    Data is inserted in chunks in order to be able to process objects that do not fit into memory.
    """

    N = len(masked_x)

    tot_max_bin_spread = 0

    for k in range(N):
        xpx = (masked_x[k]-bin_min)/delta_bin #pos of x in terms of px (from 0 to n_bins-1)
        ypx = (masked_y[k]-bin_min)/delta_bin #pos of y in terms of px
        bin_i = int(xpx) #calculates distance in i-direction from minimum bin
        bin_j = int(ypx) #calculates distance in j-direction
        Hsml = masked_h[k]/delta_bin # in number of px

        max_bin_spread = int(Hsml)
        temp =0

        kernel_norm = 3.1415926 * Hsml**2 # area of particle in px^2

        tot_max_bin_spread+=max_bin_spread
        q = masked_q[k]

        if q == 0:
            continue

        for i in range(bin_i - max_bin_spread, bin_i + max_bin_spread + 1):
            if i>=n_bins or i<0:
                continue
            for j in range(bin_j - max_bin_spread, bin_j + max_bin_spread + 1):
                if j>=n_bins or j<0:
                    continue
                u = ( (i+.5-xpx)**2 + (j+.5-ypx)**2 )**.5 / Hsml # distance of particle to px center
                if u > 1.0:
                    continue

                dx = min(xpx + Hsml, i+1) - max(xpx - Hsml, i)
                dy = min(ypx + Hsml, j+1) - max(ypx - Hsml, j)
                dxdy = dx*dy # area of px

                if Hsml < 1.:
                    wk = dxdy
                else:
                    wk = kernel(u, 1./Hsml) * dxdy

                temp += wk

                area_norm = kernel_norm / distr_weight[k] * masked_w[k] * dz[k] / delta_bin
                px_weight = area_norm * wk
                # print(f"xpx = {xpx:.3f}, ypx = {ypx:.3f},\nbin_i = {bin_i}, bin_j = {bin_j},\ndx = {dx}, dy = {dy},\nmax_bin_spread = {max_bin_spread},\ni = {i}, j = {j},\nHsml = {Hsml}, u = {u:.3f},\nwk = {wk:.3f}, temp = {temp:.6f}\npx_weight = {px_weight:.3f},\nkernel_norm = {kernel_norm:.3f},\ndistr_weight[k] = {distr_weight[k]:.6f}")
                # input()
                final_image_t[j][i] += q * px_weight
                weight_image[j][i] += px_weight
                #Implement different contributions based on bin distance (kernels)... Basic idea would be to see how large the contribution of each particle is to each pixel is
    if N>0:
        return tot_max_bin_spread/N
    else:
        return np.nan

def mapping2D(pos: np.ndarray, hsml: np.ndarray,
              qty: np.ndarray, weights: np.ndarray,
              mapParam: dict, kernel = WendlandC6_2D):
    """
    here we read input data `kv`, read chunks of particles and send them to `add_to_grid`
    """
    #set the snapshots' scalefactor and hubble factor into units.
    # units = g3u.get_units(mapParam['SNAP_PATH'])
    #produce a set of pint units from the snapshots data
    # ureg = units.get_u()

    img_xy_size = mapParam['IMG_XY_SIZE']
    img_z_size  = mapParam['IMG_Z_SIZE']
    img_pxsize = mapParam['IMG_PXSIZE'] # number of pixels per side
    img_center = np.array(mapParam['IMG_CENTER'])

    n_bins = [img_pxsize]*2
    final_image_t = np.zeros(n_bins)
    weight_image = np.zeros(n_bins)
    dw = np.zeros_like(qty)

    bins = [np.linspace(-.5,.5,img_pxsize+1), np.linspace(-.5, .5, img_pxsize+1)]

    N_part_type = len(qty)
    print(N_part_type)

    if N_part_type == 0:
        raise ValueError("Quantity array is empty... Abort")

    rel_poses = pos - img_center

    pos_x = rel_poses[:,0]
    pos_y = rel_poses[:,1]
    pos_z = rel_poses[:,2]

    # Convert into pixel position
    norm_x = (pos_x/img_xy_size)#.to('').magnitude
    norm_y = (pos_y/img_xy_size)#.to('').magnitude
    norm_z = (pos_z/img_z_size)#.to('').magnitude
    norm_h = (hsml/img_xy_size)#.to('').magnitude

    # Select indeces for cube around center based on image size
    mask = (norm_x<=.5)&(norm_x>=-.5)&(norm_y<=.5)&(norm_y>=-.5)&(norm_z<=.5)&(norm_z>=-.5)

    masked_x = norm_x[mask]
    masked_y = norm_y[mask]
    masked_h = norm_h[mask]
    masked_w = weights[mask]
    dz       = 4./3.* masked_h
    masked_q = qty[mask]#(qty[mask]).to(mapParam['RESULT_UNITS']).magnitude
    print(f"'g2D': sum of qty_arr = {np.sum(masked_q):.2e}")
    print(f"'g2D': sum of qty_arr * V_arr = {np.sum(masked_q*4/3*3.1415926*hsml[mask]**3.):.2e}")

    avg_bin_spread = calc_area_weights(dw,masked_x,masked_y,masked_h,kernel,bins[0][0], bins[0][1]-bins[0][0], img_pxsize)

    avg_bin_spread = add_to_grid(final_image_t, weight_image, masked_x, masked_y, masked_h, masked_w, masked_q, dz, kernel, dw, bins[0][0], bins[0][1]-bins[0][0], img_pxsize)
    print ('# avg bin spread ', avg_bin_spread)

    final_image = np.nan_to_num(final_image_t)

    return final_image, weight_image

def mappingParam(center: list[3], xy_size: float, Npx: int):
    mapPar = {
        "IMG_CENTER"    : center,
        "IMG_XY_SIZE"   : xy_size,
        "IMG_Z_SIZE"    : xy_size,
        "IMG_PXSIZE"    : Npx,
        "LEN_PER_PX"    : xy_size/Npx
    }
    return mapPar


from collections import defaultdict

class Solution:
    def numberOfBoomerangs(self, points):
        n = len(points)
        dist = defaultdict(lambda: 0)
        for i in range(n):
            x1, y1 = points[i]
            for j in range(i):
                x2, y2 = points[j]
                d = (x1 - x2) ** 2 + (y1 - y2) ** 2
                dist[(i, d)] += 1
                dist[(j, d)] += 1
        ans = 0
        for _, n in dist.items():
            if n > 1:
                ans += n * (n - 1)
        return ans


import art
import random
from game_data import data
from replit import clear

#FUNCIONES:
#funcion para elegir un pj random de la lista:
def random_character(data):
  character = {}
  random_index = random.randint(0, len(data)-1)
  character= data[random_index]
  data.pop(random_index)
  return character

#funcion para comparar valores y retornar el resultado (gana o pierde):
def compare_values(a_char, b_char, user_input):
  user_win = False
  a_option = 'a'
  b_option = 'b'
  if a_char['follower_count'] > b_char['follower_count']:
    if user_input == a_option:
      user_win = True
  elif a_char['follower_count'] < b_char['follower_count']:
    if user_input == b_option:
      user_win = True
  return user_win

#variables fuera del ciclo
a_char = random_character(data)
b_char = random_character(data)
score = 0
wrong_message = f"Sorry, that's wrong. Final score: {score}."
right_message = f"You're right! Current score: {score}."
empty_list = "We ran out of characters! Thank you for playing."
playing = True

#Ciclo
while playing:
  a_compare = f"Compare A: {a_char['name']} a {a_char['description']} from {a_char['country']}.\n"
  b_compare = f"Against B: {b_char['name']} a {b_char['description']} from {b_char['country']}.\n"
  print(art.logo)
  #pantalla con prints
  if score != 0:
    print(f"You're right! Current score: {score}.")
  print(a_compare)
  print(art.vs)
  print(b_compare)
  #eleccion del usuario
  user_input = input("Who has more followers? Type 'A' or 'B': ").lower()
  #se comparan los followers y devuelve el resultado
  result = compare_values(a_char, b_char, user_input)
  #resultado
  if result == True:
    score += 1
    if data == []:
      print(empty_list)
      playing = False
    clear()
    a_char = b_char
    b_char = random_character(data)
  else:
    print(f"Sorry, that's wrong. Final score: {score}.")
    playing = False


"""
It will be responsible for downloading image  captchas from our website
and save it into our disk

it will store the raw captcha images to our disk

headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36'}

"""

import requests
#the requests library will make http connections easy and heavily used in python
import argparse
import os
import time

url="https://www.e-zpassny.com/vector/jcaptcha.do"

ap=argparse.ArgumentParser()
ap.add_argument("-u","--url",help="if the images are to be downloaded from url give url",required=False,default=url)
ap.add_argument("-o","--output",required=True,help=" path to save the images")
ap.add_argument("-n","--number",type=int,default=500,help="# of images to download")
args=vars(ap.parse_args())

total=0

for i in range(0,args["number"]):
    try:
        #requesting web page from our script
        r=requests.get(args['url'],timeout=60, headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2228.0 Safari/537.3'})

        #save the image
        p = os.path.sep.join([args["output"],"{}.jpg".format(str(total).zfill(5))])
        #Str().zfill() Returns a copy of the string with '0' characters padded to the leftside of the given string.
        #ex zfill(3) ==> 000name
        f= open(p,"wb")
        f.write(r.content)
        f.close()

        print("[INFO..] Downloaded image {}".format(total))
        total+=1

    except:
        print("Error downloading image")

    time.sleep(0.2)


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

# In[27]:

import matplotlib.pyplot as plt

# In[71]:

# line plot
plt.figure(figsize=(10,3))
# cty 1
x = [2016, 2017, 2018, 2019, 2020, 2021]
y = [45, 46, 48, 50, 52, 60]
# cty 2
x2 = [2016, 2017, 2018, 2019, 2020, 2021]
y2 = [41, 42, 44, 45, 52, 57]
plt.plot(x, y, marker='o', linestyle='-.', color='g', label='mongolia')
plt.plot(x2, y2, marker='d', linestyle='-', color='b', label='congo')
plt.xlabel('Year')
plt.ylabel('Population (M)')
plt.title('years vs population')
plt.legend(loc='upper center')
plt.yticks([45, 50, 55, 60])
plt.savefig('cty.png')
plt.show()

# In[76]:

#subplots
fig, ax = plt.subplots(1, 2, sharey=True) # rows and cols
ax[0].plot(x,y,color='r')
ax[1].plot(x2,y2,color='b')
plt.show()

# In[70]:

# Bar Plots
plt.figure(figsize=(8,3))
bar1 = ['asia', 'america', 'europe', 'oceania', 'africa']
bar2 = [30, 20, 20, 5, 25]
plt.bar(bar1, bar2)
plt.show()

# In[48]:

# pie
plt.pie(y, labels=x)
plt.show()

# In[57]:

# histogram
conts = ['asia', 'america', 'europe', 'oceania', 'africa']
pers = [30, 15, 25, 5, 25]
bins = [5,10,15,20,25,30,35]
plt.hist(pers,bins,edgecolor='black')
plt.show()

# In[69]:

# scatterplot
conts = ['asia', 'america', 'europe', 'oceania', 'africa']
scat1 = [30,15,25,5,25,50]
scat2 = [5,10,12,2,1,4]
plt.scatter(scat1,scat2)
plt.show()

# In[ ]:


import numpy as np
import numpy.random as rd
import matplotlib.pyplot as plt
import seaborn as sns  # noqa
from itertools import islice

from ibm_rng import ibm_rng

from pylatex import Plt

N = 1000

def main():
    dist_to_tex('np_random.tex', 'Numpy RNG', *get_random_pairs(N).T)

    ibm_gen = ibm_rng(1)
    dist_to_tex('ibm_random.tex', 'IBM RNG',
                take_samples(ibm_gen), take_samples(ibm_gen))

def take_samples(gen, n=N):
    return np.array(list(islice(gen, n)))

def dist_to_tex(filename, caption, x, y):
    plot_distribution(x, y, show=False)
    with open(filename, 'w') as f:
        plot = Plt(position='htbp')
        plot.add_plot(plt)
        plot.add_caption(caption)

        plot.dump(f)

def plot_distribution(x, y, show=True):
    sns.jointplot(x, y, stat_func=None)
    if show:
        plt.show()

def get_random_pairs(n):
    return rd.rand(n, 2)

if __name__ == '__main__':
    main()


import os
import json
import tensorflow as tf
from keras.models import Sequential
import keras
from keras.layers import Dense, Dropout, LSTM, CuDNNLSTM, Activation, BatchNormalization, LeakyReLU, Flatten, TimeDistributed, CuDNNGRU, RNN, SimpleRNN, GRU
import numpy as np
import random
from sklearn.preprocessing import normalize
from tensorflow.keras.callbacks import TensorBoard
from keras import backend as K
import time
from normalizer import norm
import matplotlib.pyplot as plt
import pickle
import sys
from sklearn import preprocessing
import itertools
from sklearn.model_selection import train_test_split

#gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
#sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

class Visualize(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs={}):
        if epoch % 10 == 0 or 1==1:
            preds = []
            grounds = []
            x = list(range(len(test_samples)))
            for i in test_samples:
                preds.append(model.predict([[x_test[i]]])[0][0])
                grounds.append(y_test[i])

            plt.clf()
            plt.plot(x, grounds, label='ground-truth')#, marker='o')
            plt.plot(x, preds, label='prediction-{}'.format(epoch))#, marker='o')
            plt.title("test data (unseen)")
            plt.legend()
            plt.savefig("models/h5_models/{}/0-{}-epoch-{}.png".format(model_name, model_name, epoch))
            plt.pause(.1)
            model.save("models/h5_models/{}/{}-epoch-{}.h5".format(model_name, model_name, epoch))

#v_ego, v_lead, d_lead
data_dir = "LSTM"
os.chdir("C:/Git/dynamic-follow-tf")

norm_dir = "data/{}/normalized.npy"

model_name = "LSTM"

'''with open("data/{}/x_train".format(data_dir), "rb") as f:
    x_train = pickle.load(f)

with open("data/{}/y_train".format(data_dir), "rb") as f:
    y_train = pickle.load(f)'''

samples_to_use = 7000000
if samples_to_use != 'all':
    y_train = np.load("data/{}/y_train.npy".format(data_dir))[:samples_to_use]
else:
    y_train = np.load("data/{}/y_train.npy".format(data_dir))
    print(len(y_train))
is_array = False

if not os.path.exists(norm_dir.format(data_dir)):
    if samples_to_use != 'all':
        x_train = np.load("data/{}/x_train.npy".format(data_dir))[:samples_to_use]
    else:
        x_train = np.load("data/{}/x_train.npy".format(data_dir))
    print("Normalizing...", flush=True)
    normalized = norm(x_train)
    x_train = normalized['normalized']
    scales = normalized['scales']
    print("Dumping normalization...", flush=True)
    np.save(norm_dir.format(data_dir), x_train)
    with open('data/LSTM/scales', "wb") as f:
        pickle.dump(normalized['scales'], f)
    #with open(norm_dir.format(data_dir), "wb") as f:
        #pickle.dump(normalized, f)
else:
    is_array = True
    print("Loading normalized data...", flush=True)
    x_train = np.load(norm_dir.format(data_dir))
    with open('data/LSTM/scales', "rb") as f:
        scales = pickle.load(f)
    print('Loaded!', flush=True)
    #with open(norm_dir.format(data_dir), "rb") as f:
        #normalized = pickle.load(f)

print(len(x_train))
#scales = normalized['scales']
#x_train = normalized['normalized']
y_train = np.array([np.interp(i, [-1, 1], [0, 1]) for i in y_train])

flatten = True
if flatten:
    print('Flattening...', flush=True)
    #x_train = np.array([[inner for outer in sample for inner in outer] for sample in x_train])
    x_train = np.array([i.flatten() for i in x_train])  # whole lot faster lul
elif not is_array:
    x_train = np.array(x_train)
print(x_train.shape)

x_train, x_test, y_train, y_test = train_test_split(x_train, y_train, test_size=0.1)

#random_choices = []
num_test = 600
num_test = min(len(x_test), num_test)

test_samples = random.sample(range(len(x_test)), num_test)
test_samples = sorted(list(zip([y_test[i] for i in test_samples], test_samples)))
test_samples = [y for x, y in test_samples] # sort from lowest to highest for visualization

opt = keras.optimizers.Adam(lr=0.0001)#, decay=1e-6)
opt = keras.optimizers.Adadelta()
#opt = keras.optimizers.RMSprop(0.001)
#opt = keras.optimizers.Adagrad(lr=0.00001)
opt = 'adam'
#opt = 'rmsprop'
#opt = keras.optimizers.SGD(lr=0.01, decay=1e-8, momentum=0.9, nesterov=True)

layers = 6
nodes = 186

model_type = 'dense' # lstm, dense, or gru

model = Sequential()
if model_type.lower() == 'lstm':
    to_sub = 2
    model.add(CuDNNLSTM(nodes, input_shape=(x_train.shape[1:]), return_sequences=True))
elif model_type.lower() == 'dense':
    to_sub = 1
    model.add(Dense(nodes, activation="relu", input_shape=(x_train.shape[1:])))
    #model.add(Dropout(0.1))
#model.add(CuDNNGRU(nodes, return_sequences=True, input_shape=(x_train.shape[1:])))
#model.add(SimpleRNN(128, activation='tanh', return_sequences=False, input_shape=(x_train.shape[1:])))
#model.add(Dropout(.05))
for i in range(layers - to_sub):
    #model.add(Dense(nodes, activation="relu"))
    #model.add(Dropout(.05))
    if model_type.lower() == 'lstm':
        model.add(CuDNNLSTM(nodes, return_sequences=True))
    #model.add(CuDNNGRU(nodes, return_sequences=True))
    elif model_type.lower() == 'dense':
        model.add(Dense(nodes, activation="relu"))
        #model.add(Dropout(0.1))
    #model.add(SimpleRNN(64, activation='tanh'))
#model.add(Dense(128, activation='relu'))
#model.add(Permute((2,1), input_shape=(10, 5)))
#model.add(CuDNNGRU(nodes, return_sequences=False))
if model_type.lower() == 'lstm':
    model.add(CuDNNLSTM(nodes, return_sequences=False))

model.add(Dense(1, activation='linear'))

model.compile(loss='mean_squared_error', optimizer=opt)
#tensorboard = TensorBoard(log_dir="logs/test-{}".format("30epoch"))
callback_list = [Visualize()]
model.fit(x_train, y_train, shuffle=True, batch_size=256, epochs=200, validation_data=(x_test, y_test), callbacks=callback_list)

def get_acc():
    accs = []
    for idx, i in enumerate(x_test[:20000]):
        pred = model.predict([[i]])[0][0]
        accs.append(abs(pred - y_test[idx]))
    print('Test accuracy: {}'.format(1 - sum(accs) / len(accs)))

    '''accs = []
    for idx, i in enumerate(x_train):
        pred = model.predict([[i]])[0][0]
        accs.append(abs(pred - y_train[idx]))
    print('Train accuracy: {}'.format(1 - sum(accs) / len(accs)))'''

get_acc()

for i in range(10):
    rand = random.randint(0, len(x_train))
    pred = model.predict([[x_train[rand]]])[0][0]
    print("Ground truth: {}".format(y_train[rand]))
    print("Prediction: {}\n".format(pred))

save_model = False
if save_model:
    model.save("models/h5_models/"+model_name+".h5")
    print("Saved model!")


'''
Module to define transformations on dataframes
'''
import pyspark.sql.functions as F

def transform_data(input_df):
    ''' This function takes a dataframe as input and returns a transformed
    dataframe'''
    transformed_df = (input_df
                      .groupBy('Location',)
                      .agg(F.sum('ItemCount').alias('TotalItemCount')))
    return transformed_df


import pytest
from removeDuplicates import Solution

@pytest.mark.parametrize("nums,expected", [
    ([1, 2, 2, 3], [1, 2, 2, 3]),
    ([1, 1, 1, 2, 2, 3], [1, 1, 2, 2, 3]),
    ([1, 1, 1, 2, 2, 3, 3, 3, 3], [1, 1, 2, 2, 3, 3]),
    ([0, 0, 1, 1, 1, 1, 2, 3, 3], [0, 0, 1, 1, 2, 3, 3])
])
def test_removeDuplicates(nums, expected):
    length = Solution().removeDuplicates(nums)
    assert nums == expected
    assert len(nums) == length


"""The services tests module."""
import pytest
from django.conf import settings
from django.core.exceptions import ValidationError
from django.db.models.query import QuerySet

from communication.models import Review, ReviewComment
from communication.validators import (validate_message_parent,
                                      validate_message_recipient,
                                      validate_reminder_remind_before,
                                      validate_review_comment_user,
                                      validate_review_user)
from users.models import User

pytestmark = pytest.mark.django_db

def test_validate_reminder_remind_before():
    """Should validate remind_before field."""
    step: int = settings.D8B_REMINDER_INTERVAL
    with pytest.raises(ValidationError):
        validate_reminder_remind_before(int(step * 1.5))
    validate_reminder_remind_before(step * 3)

def test_validate_review_comment_user(admin: User, reviews: QuerySet):
    """Should validate the review comment user."""
    comment = ReviewComment()
    comment.user = admin
    comment.review = reviews.exclude(professional__user=admin).first()

    with pytest.raises(ValidationError):
        validate_review_comment_user(comment)

def test_validate_review_user(admin: User, professionals: QuerySet):
    """Should validate the review user."""
    review = Review()
    review.user = admin
    review.professional = professionals.filter(user=admin).first()
    review.rating = 3  # type: ignore
    review.description = "description"

    with pytest.raises(ValidationError):
        validate_review_user(review)

def test_validate_message_recipient(admin: User, messages: QuerySet):
    """Should mark the message read."""
    message = messages.filter(sender=admin).first()
    validate_message_recipient(message)
    message.recipient = admin
    with pytest.raises(ValidationError):
        validate_message_recipient(message)

def test_validate_parent(admin: User, messages: QuerySet):
    """Should validate the message parent message."""
    messages = messages.filter(sender=admin)
    message = messages[0]
    message.parent = messages[1]
    with pytest.raises(ValidationError):
        validate_message_parent(message)
    message.is_read = True
    with pytest.raises(ValidationError):
        validate_message_parent(message)
    message.parent = messages.filter(recipient=admin).first()
    validate_message_parent(message)


"""Module to describe an atomic structure."""
import json
from jarvis.db.figshare import get_jid_data
from jarvis.core.atoms import Atoms
from jarvis.analysis.structure.spacegroup import Spacegroup3D
from jarvis.analysis.diffraction.xrd import XRD
from jarvis.core.specie import Specie
import pprint
from collections import defaultdict

def atoms_describer(
    atoms=[], xrd_peaks=5, xrd_round=1, cutoff=4, take_n_bomds=2
):
    """Describe an atomic structure."""
    spg = Spacegroup3D(atoms)
    theta, d_hkls, intens = XRD().simulate(atoms=(atoms))
    #     x = atoms.atomwise_angle_and_radial_distribution()
    #     bond_distances = {}
    #     for i, j in x[-1]["different_bond"].items():
    #         bond_distances[i.replace("_", "-")] = ", ".join(
    #             map(str, (sorted(list(set([round(jj, 2) for jj in j])))))
    #         )
    dists = defaultdict(list)
    elements = atoms.elements
    for i in atoms.get_all_neighbors(r=cutoff):
        for j in i:
            key = "-".join(sorted([elements[j[0]], elements[j[1]]]))
            dists[key].append(j[2])
    bond_distances = {}
    for i, j in dists.items():
        dist = sorted(set([round(k, 2) for k in j]))
        if len(dist) >= take_n_bomds:
            dist = dist[0:take_n_bomds]
        bond_distances[i] = ", ".join(map(str, dist))
    fracs = {}
    for i, j in (atoms.composition.atomic_fraction).items():
        fracs[i] = round(j, 3)
    info = {}
    chem_info = {
        "atomic_formula": atoms.composition.reduced_formula,
        "prototype": atoms.composition.prototype,
        "molecular_weight": round(atoms.composition.weight / 2, 2),
        "atomic_fraction": json.dumps(fracs),
        "atomic_X": ", ".join(
            map(str, [Specie(s).X for s in atoms.uniq_species])
        ),
        "atomic_Z": ", ".join(
            map(str, [Specie(s).Z for s in atoms.uniq_species])
        ),
    }
    struct_info = {
        "lattice_parameters": ", ".join(
            map(str, [round(j, 2) for j in atoms.lattice.abc])
        ),
        "lattice_angles": ", ".join(
            map(str, [round(j, 2) for j in atoms.lattice.angles])
        ),
        "spg_number": spg.space_group_number,
        "spg_symbol": spg.space_group_symbol,
        "top_k_xrd_peaks": ", ".join(
            map(
                str,
                sorted(list(set([round(i, xrd_round) for i in theta])))[
                    0:xrd_peaks
                ],
            )
        ),
        "density": round(atoms.density, 3),
        "crystal_system": spg.crystal_system,
        "point_group": spg.point_group_symbol,
        "wyckoff": ", ".join(list(set(spg._dataset["wyckoffs"]))),
        "bond_distances": bond_distances,
        "natoms_primitive": spg.primitive_atoms.num_atoms,
        "natoms_conventional": spg.conventional_standard_structure.num_atoms,
    }
    info["chemical_info"] = chem_info
    info["structure_info"] = struct_info
    return info

if __name__ == "__main__":
    atoms = Atoms.from_dict(
        get_jid_data(jid="JVASP-32", dataset="dft_3d")["atoms"]
    )
    info = atoms_describer(atoms=atoms)
    pprint.pprint(info)
    from pygments import highlight
    from pygments.formatters.terminal256 import Terminal256Formatter
    from pygments.lexers.web import JsonLexer

    raw_json = json.dumps(info, indent=4)
    colorful = highlight(
        raw_json,
        lexer=JsonLexer(),
        formatter=Terminal256Formatter(),
    )
    print(colorful)