#
# Little Script for Geotag Pictures of a Garmin DashCam
# Robert Benet - 30/08/2020
#
# Information extract from:
# How to configur Tesseract
# https://stackoverflow.com/questions/50951955/pytesseract-tesseractnotfound-error-tesseract-is-not-installed-or-its-not-i/53672281
# Install Teseract for windows
# https://github.com/UB-Mannheim/tesseract/wiki
# Install Tesseact for Python
# https://pypi.org/project/pytesseract/
# Tutorial for Tesseract
# https://nanonets.com/blog/ocr-with-tesseract/
# EXIF
# https://pypi.org/project/exif/
# pip install exif
# pyexiv2 - No works fine
# https://pypi.org/project/pyexiv2/
# OpenCV
# pip install opencv-python

import pytesseract
from pytesseract import Output
import cv2
import numpy as np
from exif import Image
#from pyexiv2 import Image

class Ocr_Dascam():
    def __init__(self,dir_to_find=None,oem=None,psm=None,parent=None,filter_image_selected='Dilate'):
        from os import scandir,path

        # App Message
        if parent==None:
            self.parent=self
            self.parent.info=self.info
        else:
            self.parent=parent

        #check exist and set path of EXE Tesseract
        exe_tesseract=r'C:\Program Files\Tesseract-OCR\tesseract.exe'
        if not path.exists(exe_tesseract):
            self.parent.info('File not exist: '+exe_tesseract)
            return
        pytesseract.pytesseract.tesseract_cmd = exe_tesseract

        # Adding custom options
        if not oem==None and not psm==None:
            custom_config=r'--oem '+str(oem)+' --psm '+str(psm[:2])
        elif not oem==None:
            custom_config=r'--oem '+str(oem)
            custom_config=r' --psm'+str(psm)
        else:
            custom_config = r'--oem 3 '

        # Select Photo Directory
        if dir_to_find==None:
            self.select_directory()
        else:
            self.dir_to_find=dir_to_find

        result={}
        correct=0
        incorrect=0
        # Create TXT File with Results
        file_txt=open(path.join(self.dir_to_find,'Result_OCR_Scan_'+filter_image_selected+'.txt'),'w')
        file_txt.write('Configuration '+custom_config+chr(10))
        file_txt.write('Filter '+filter_image_selected+chr(10)+chr(10))

        for file in scandir(self.dir_to_find):
            # extract name file and extension separate
            name,ext=path.splitext(file.name)

            # Find coordinates from a picture file
            if ext.upper()=='.JPG' or ext.upper()=='.JPEG' or ext.upper()=='.PNG':
                # Update date Creation of JPG
                capture_time_sec_crea=path.getctime(file) # Creation Date
                capture_time_sec_modi=path.getmtime(file) # Modification Date
                if capture_time_sec_crea<capture_time_sec_modi: # Sometimes creation date is wrong
                    capture_time_sec=capture_time_sec_crea
                else:
                    capture_time_sec=capture_time_sec_modi

                find=False
                incorrect+=1
                day=' No Encontrado '
                hour=''
                lat=''
                lon=''

                # Filter Image for a good reading
                self.filter_image(file.path,filter_image_selected)

                read = pytesseract.image_to_data(self.img, output_type=Output.DICT)
                day=' No Encontrado '+'.'.join(str(x).strip() for x in read['text'])
                #print(file.path,read['text'])

                for i in range(len(read['text'])):
                    if read['text'][i]=='GARMIN':
                        if len(read['text'])>i+4:

                            day=read['text'][i+1]
                            hour=read['text'][i+2]
                            try:
                                lat=float(read['text'][i+3])
                            except ValueError:
                                break
                            try:
                                lon=float(read['text'][i+4])
                            except ValueError:
                                break
                            if lat>-90 and lat<90 and lon>-180 and lon<180 and not str(lat).strip()=='' and not str(lon).strip()=='':
                                self.geotag_photo(file.path,lat,lon)
                                self.changeFileCreationTime(file.path,capture_time_sec)
                                find=True
                                incorrect-=1
                                correct+=1
                            break

                result[file.path]={'day':day,'hour':hour,'lat':lat,'lon':lon,'find':find}
                text_result=file.path+", "+day+", "+hour+",    GPS:,"+str(lat)+","+str(lon)+","+str(find)
                file_txt.write(text_result+chr(10))
                self.parent.info(text_result)

        #self.parent.info(result)
        self.parent.info('GeoLocalizados   : '+str(correct))
        self.parent.info('No GeoLocalizados: '+str(incorrect))

        file_txt.write('GeoLocalizados   : '+str(correct)+chr(10))
        file_txt.write('No GeoLocalizados: '+str(incorrect)+chr(10))
        file_txt.close()

    # Change File Creation
    def changeFileCreationTime(self,fname, newtime):
        # https://stackoverflow.com/questions/4996405/how-do-i-change-the-file-creation-date-of-a-windows-file
        import pywintypes, win32file, win32con
        wintime = pywintypes.Time(newtime)
        winfile = win32file.CreateFile(
            fname, win32con.GENERIC_WRITE,
            win32con.FILE_SHARE_READ | win32con.FILE_SHARE_WRITE | win32con.FILE_SHARE_DELETE,
            None, win32con.OPEN_EXISTING,
            win32con.FILE_ATTRIBUTE_NORMAL, None)

        win32file.SetFileTime(winfile, wintime, None, None)
        winfile.close()

    # Geotag Photo
    def geotag_photo(self,file,lat,lon):

        exif_lat=(int(lat),int(lat % 1 * 60),(lat % 1 * 60)%1*60)
        exif_lon=(int(lon),int(lon % 1 * 60),(lon % 1 * 60)%1*60)
        if int(lat)<0:
            exif_lat_ref='S'
        else:
            exif_lat_ref='N'

        if int(lon)<0:
            exif_lon_ref='W'
        else:
            exif_lon_ref='E'

        with open(file, 'rb') as image_file:
            my_image = Image(image_file)
        my_image.make = "Python"
        my_image.gps_latitude_ref=exif_lat_ref
        my_image.gps_latitude=exif_lat
        my_image.gps_longitude_ref= exif_lon_ref
        my_image.gps_longitude= exif_lon

        with open(file, 'wb') as new_image_file:
            new_image_file.write(my_image.get_file())

        #self.parent.info(exif_lat,"_",exif_lon)

    # Select directory
    def select_directory(self):
        from tkinter import filedialog
        from tkinter import Tk
        root = Tk()
        root.withdraw()
        self.dir_to_find = filedialog.askdirectory()

    # App Message
    def info(self,message):
        print (message)

    def filter_image(self,file,filter):
        img = cv2.imread(file)
        img=self.get_grayscale(img)
        if filter==None or filter=='Dilate':
            self.img=self.dilate(img)
        elif filter=='Thresholding':
            self.img=self.thresholding(img)
        elif filter=='Remove Noise':
            self.img=self.remove_noise(img)
        elif filter=='Erosion':
            self.img=self.erode(img)
        elif filter=='Opening':
            self.img=self.opening(img)
        elif filter=='Canny':
            self.img=self.canny(img)
        elif filter=='Eskew':
            self.img=self.deskew(img)

    # gray scale
    def get_grayscale(self,image):
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # noise removal
    def remove_noise(self,image):
        return cv2.medianBlur(image,5)

    #thresholding
    def thresholding(self,image):
        return cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]

    #dilation
    def dilate(self,image):
        kernel = np.ones((5,5),np.uint8)
        return cv2.dilate(image, kernel, iterations = 1)

    #erosion
    def erode(self,image):
        kernel = np.ones((5,5),np.uint8)
        return cv2.erode(image, kernel, iterations = 1)

    #opening - erosion followed by dilation
    def opening(self,image):
        kernel = np.ones((5,5),np.uint8)
        return cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)

    #canny edge detection
    def canny(self,image):
        return cv2.Canny(image, 100, 200)

    #skew correction
    def deskew(self,image):
        coords = np.column_stack(np.where(image > 0))
        angle = cv2.minAreaRect(coords)[-1]
        if angle < -45:
            angle = -(90 + angle)
        else:
            angle = -angle
        (h, w) = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        return rotated

    #template matching
    def match_template(self,image, template):
        return cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)

if __name__ == "__main__":
    # Load Main Modul
    #action=Ocr_Dascam()

    import tkinter as tk
    from tkinter import ttk

    app = tk.Tk()
    app.title("GeoTag Photos from Garmin Dash Cam")
    app.geometry('500x200')

    line0=tk.Label(app, text = '')
    line0.grid(column=0, row=0)

    labelOEM = tk.Label(app, text = "OEM")
    labelOEM.grid(column=0, row=2,sticky='W',padx=5)

    comboOEM = ttk.Combobox(app,values=[0,1,2,3])
    comboOEM.grid(column=1, row=2,sticky='W',padx=5)
    comboOEM.current(3)

    labelPSM = tk.Label(app, text = "PSM")
    labelPSM.grid(column=0, row=3,sticky='W',padx=5)

    comboPSM= ttk.Combobox(app, width=55,values=['0 Orientation and script detection (OSD) only.','1    Automatic page segmentation with OSD.','2    Automatic page segmentation, but no OSD, or OCR.','3    Fully automatic page segmentation, but no OSD. (Default)','4    Assume a single column of text of variable sizes.','5    Assume a single uniform block of vertically aligned text.','6    Assume a single uniform block of text.','7    Treat the image as a single text line.','8    Treat the image as a single word.','9    Treat the image as a single word in a circle.','10    Treat the image as a single character.','11    Sparse text. Find as much text as possible in no particular order.','12    Sparse text with OSD.','13    Raw line. Treat the image as a single text line, bypassing hacks that are'])
    comboPSM.grid(column=1, row=3,sticky='W',padx=5)
    comboPSM.current(3)

    labelFilter = tk.Label(app, text = "Photo Filter")
    labelFilter.grid(column=0, row=5,sticky='W',padx=5)

    comboFilter= ttk.Combobox(app,values=['Dilate','Thresholding','Remove Noise','Erosion','Opening','Canny','Eskew'])
    comboFilter.grid(column=1, row=5,sticky='W',padx=5)
    comboFilter.current(0)

    line6=tk.Label(app, text = '')
    line6.grid(column=0, row=6)

    buton_action=tk.Button(app, text = "Execute",command=lambda: Ocr_Dascam(oem=comboOEM.get(),psm=comboPSM.get(),filter_image_selected=comboFilter.get()))
    buton_action.grid(column=1, row=7,sticky='E',padx=5)

    app.mainloop()
    app.quit()

    #print(comboOEM .current(), comboOEM.get())
    #action=Ocr_Dascam(dir_to_find=r'C:\GisBike\Test\SECRETARIA\VOLTA 2020\RECORREGUTS\ETAPA 1\FOTOS\RECORREGUT_TEST')


import asyncio
import atexit
import json
import random
import re
import secrets
import time
from datetime import datetime
from typing import List

import nest_asyncio
import uvicorn
import websockets
from azure.iot.device.aio import IoTHubDeviceClient
from fastapi import (
    Depends,
    FastAPI,
    HTTPException,
    Request,
    Response,
    WebSocket,
    status,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from sqlalchemy.orm import Session
from sqlalchemy.sql.functions import user

from api import crud, models
from database import schemas
from database.database import SessionLocal, engine
from helper import get_dispense_duration, get_pump_id, running_on_raspberrypi
from services import barrels as barrels_service
from services import iotdevices
from services import pump as pump_service
from services import rfid, senddispenserinfo

nest_asyncio.apply()

# async def handler(websocket, path):
#     while True:
#         id = random.randint(1, 2)
#         name = "John Doe" if id == 1 else "Fred"
#         data = [{"id": id, "name": name}]
#         await asyncio.sleep(30)
#         await websocket.send(json.dumps(data))

# start_server = websockets.serve(handler, "127.0.0.1", 8888)
# asyncio.get_event_loop().run_until_complete(start_server)

models.Base.metadata.create_all(bind=engine)
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def cleanup():
    print(f"[{datetime.utcnow()}] Process exit.")
    if running_on_raspberrypi():
        pump_service.cleanup()

atexit.register(cleanup)

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

@app.post("/userInfo/", response_model=schemas.User, tags=["user"])
async def get_user(request: Request, db: Session = Depends(get_db)):
    body = await request.json()
    db_user = crud.get_user_by_name(body["name"], db=db)

    if db_user == None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="User not found."
        )
    return db_user

@app.post("/dispenses/", tags=["dispense"])
async def create_dispense(request: Request, db: Session = Depends(get_db)):
    body = await request.json()
    print(f"[{datetime.utcnow()}] Start dispensing oil: ", body)

    carManufacturer = (
        body["carManufacturer"] if body["carManufacturer"] != "None" else ""
    )
    carModel = body["carModel"] if body["carModel"] != "None" else ""
    carEngineType = body["carEngineType"] if body["carEngineType"] != "None" else ""

    if running_on_raspberrypi():
        pump_id = get_pump_id(body["oilType"])
        dispense_duration = get_dispense_duration(body["dispensedQuantity"])
        pump_service.control_pump("start", pump_id)
        time.sleep(dispense_duration)
        pump_service.control_pump("stop", pump_id)

    dispense = models.Dispense(
        car_registration_number=body["registrationNumber"],
        car_manufacturer=carManufacturer,
        car_model=carModel,
        car_engine_type=carEngineType,
        suggested_quantity=body["suggestedQuantity"],
        dispensed_quantity=body["dispensedQuantity"],
        dispense_type=body["dispenseType"],
        user_id=body["userId"],
        created_at=datetime.utcnow(),
        barrels_id=barrels_service.get_barrel_id(body["oilType"], db=db),
    )

    db_dispense = crud.create_dispense(db=db, dispense=dispense)
    print("dispense results", db_dispense.id)

    crud.update_barrels(
        id=barrels_service.get_barrel_id(body["oilType"], db=db),
        dispensed_quantity=body["dispensedQuantity"],
        db=db,
    )
    print(
        f"barrels table updated for id: {barrels_service.get_barrel_id(body['oilType'],db=db)}"
    )

    get_remaining_quantity = crud.get_remaining_quantity(
        id=barrels_service.get_barrel_id(body["oilType"], db=db), db=db
    )

    data = {}
    for results in get_remaining_quantity:
        data = {
            "data": {
                "oil_type": results[0],
                "remaining_quantity": results[1],
                "remaining_quantity_perc": round(results[2], 2),
                "remainingInBarrel": "73.22 L (35.85%)",
            }
        }

    cloud_results = await crud.get_cloud_dispense_data(db_dispense.id, db=db)
    print("cloud results", cloud_results)

    send_info_to_azure = senddispenserinfo.SendDispensingIinfo(
        iotdevices.iot_devices["rd01"]
    )

    if cloud_results == None:
        print("nothing returned")
        pass
    else:
        await send_info_to_azure.send_dispense_info(
            cloud_results[0][0],
            cloud_results[0][1].strftime("%Y-%m-%d, %H:%M:%S"),
            cloud_results[0][2],
            cloud_results[0][3],
            cloud_results[0][4],
            cloud_results[0][5],
            cloud_results[0][7],
            cloud_results[0][6],
            cloud_results[0][8],
            cloud_results[0][9],
            cloud_results[0][10],
            data["data"]["remaining_quantity"],
            data["data"]["remaining_quantity_perc"],
        )

    return data

@app.get("/dispenses/{user_id}", tags=["dispense"])
def get_dispense(user_id: int, db: Session = Depends(get_db)):
    get_transaction = crud.get_dispense(user_id=int(user_id), db=db)
    return get_transaction

@app.get("/barrel/", tags=["barrel"])
def get_barrel(db: Session = Depends(get_db)):
    db_barrel = crud.get_barrel(db)
    return db_barrel

@app.put("/barrel/{id}", tags=["barrel"])
def get_barrel(id: int, db: Session = Depends(get_db)):
    db_barrel = crud.update_barrels_by_id(id=id, db=db)
    return db_barrel

@app.get("/barreltypes/", response_model=List[schemas.BarrelType], tags=["barreltype"])
def get_barrel_types(db: Session = Depends(get_db)):
    db_barrel_types = crud.get_barrel_types(db)
    return db_barrel_types

@app.post("/order/", response_model=schemas.Order, tags=["order"])
async def create_order(order: schemas.OrderCreate, db: Session = Depends(get_db)):
    crud.create_order(db=db, order=order)
    data = """{success}"""
    return Response(content=data, media_type="application/xml")

@app.get("/order/", tags=["order"])
def get_order(db: Session = Depends(get_db)):
    db_orders = crud.get_order_records(db=db)
    return db_orders

@app.get("/orders/{user_id}", tags=["order"])
def get_order_transactions(user_id: int, db: Session = Depends(get_db)):
    get_transactions = crud.get_order_transactions(user_id=int(user_id), db=db)
    return get_transactions

@app.get("/environmental-benefits/{user_id}", tags=["order"])
def get_environmental_benefits(user_id: int, db: Session = Depends(get_db)):
    transactions = crud.get_environmental_benefits(user_id=int(user_id), db=db)

    result = []

    for t in transactions:
        enhanced_item = {
            "id": t[0],
            "created_at": t[1],
            "dispensed_quantity": t[2],
            "saved_plastic": t[3],
            "saved_co2": t[4],
            "packs": 1 if t[2] < 3 else 2,
        }
        result.append(enhanced_item)

    return result

@app.get("/oilTypes/reg/", tags=["external db"])
async def get_oil_types_by_reg_number(vehicle_reg: str, db: Session = Depends(get_db)):
    print(f"[{datetime.utcnow()}] Received registration number: {vehicle_reg}.")

    MIN_NUMBER_LENGTH = 6

    if len(vehicle_reg) < MIN_NUMBER_LENGTH:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Car Registration Number not found.",
        )
    else:
        return {
            "recommendedQuantity": "4.5",
            "details": "Volkswagen (VW) (EU) Passat B7, 3C/36 (2010-2015)",
            "types": crud.get_barrel_records(db),
        }

@app.get("/oilTypes/", tags=["external db"])
def get_oil_types(db: Session = Depends(get_db)):
    print(f"[{datetime.utcnow()}] Received request to get the oil types.")
    db_barrel_records = crud.get_barrel_records(db)
    return db_barrel_records

@app.post("/api/pump/procedure", tags=["pump"])
async def pump_start(request: Request):
    body = await request.json()
    print(f"[{datetime.utcnow()}] Received command for Pump {body['pumpId']}.", body)

    if running_on_raspberrypi():
        pump_service.control_pump(body["action"], body["pumpId"])

    return

##RFID Login Handling - Begin

# Start the RFID reader to keep reading forever
rfid.start_rfid_reader()

# set up the rfid_read_event to notify when a value has been read from the RFID
global rfid_read_event
rfid_read_event = asyncio.Event()

# Create and start the Websocket for RFID login information to be sent to client
@app.websocket("/rfid_login")
async def rfid_login_handler(websocket: WebSocket):
    global rfid_user
    await websocket.accept()
    while True:
        await rfid_read_event.wait()
        if rfid_user is not None:
            user_json = json.dumps(
                {"id": rfid_user["id"], "name": f"{rfid_user['name']}"}
            )
            print(f"Notifying client of RFID User for {user_json}")
            await websocket.send_text(user_json)
            rfid_user = None
        rfid_read_event.clear()

# Given an RFID token id, looks up user in database and sends notification to front end that user wants to log in
def notify_client(rfid_token_id, db):
    global rfid_user
    rfid_user = crud.get_user_by_rfid_tag(rfid_token_id, db)
    print(f"Found User for rfid_token_id {rfid_token_id}: {rfid_user}")
    rfid_read_event.set()

##RFID Login Handling - End


import pandas as pd
import numpy as np
from enum import Enum
from datetime import time

class SleepValue(Enum):
    none = 0
    sleeping = 1
    restless = 2
    awake = 3

SLEEP_VALUES = {x.value:x.name for x in SleepValue}

STATS_NAME_BASIC = 'basicStats'
STATS_NAME_TIMING = 'timingStats'
STATS_NAME_BASIC_AND_TIMING = 'basicAndTimingStats'
STATS_NAME_INTERVALS = 'intervalsStats'
STATS_NAME_INTRADAY = 'intradayStats'

def generateStatsFrom(data, statsType, **kwargs):
    """
    Generic method for stats generation.
    :param data: sleep data. List of dataframes, one for each day
    :param statsType: name of the stat to be generated
    :return:
    """

    if statsType == STATS_NAME_BASIC:
        stats = generateBasicStats(data)
    elif statsType == STATS_NAME_TIMING:
        stats = generateTimingStats(data, **kwargs)
    elif statsType == STATS_NAME_BASIC_AND_TIMING:
        basicStats = generateBasicStats(data)
        timingStats = generateTimingStats(data, **kwargs)
        stats = pd.concat([basicStats, timingStats], axis=1)
    elif statsType == STATS_NAME_INTERVALS:
        stats = generateIntervalsStats(data, **kwargs)
    elif statsType == STATS_NAME_INTRADAY:
        stats = generateIntradayStats(data, **kwargs)
    else:
        raise Exception(statsType + ' Stat not implemented')
    return stats

#Each row is a day with basic stats like sleep values count, efficiency and total minutes
def generateBasicStats(filesData):
    #Create new df to store the stats
    basicStats = pd.DataFrame(columns=[1,2,3,'total_minutes','sleep_efficiency'])
    basicStats.index.rename('date', inplace=True)

    #For each file-data, extract the basic stats, and add them to the df
    for fileData in filesData:
        stats = fileData.groupby(['value'], as_index=False).size()
        stats['total_minutes'] = stats.sum()
        date = getSleepDate(fileData)
        basicStats.loc[date] = stats

    #Derive additional stats
    basicStats = basicStats.fillna(0).astype(int)
    basicStats['sleep_efficiency'] = (basicStats[SleepValue.sleeping.value]/basicStats['total_minutes'])*100
    basicStats['sleep_inefficiency'] = (basicStats['sleep_efficiency']-100).abs()
    basicStats['sleep_hours'] = (basicStats[SleepValue.sleeping.value]/60)
    basicStats['total_hours'] = basicStats['total_minutes']/60

    #Rename columns
    columns = SLEEP_VALUES
    columns['value'] = 'date'
    basicStats.rename(columns=columns, inplace=True)
    return basicStats

#Each row is a day with timing stats about like first minute asleep, interval avg and max length
def generateTimingStats(filesData, minSleepIntervalRequired=0):
    #Create new df to store the stats
    intervalsStats = pd.DataFrame(columns=['first_min_asleep',
                                           'to_bed_time',
                                           'wake_up_time',
                                           'sleep_interval_max_len',
                                           'sleep_interval_avg_len'])
    intervalsStats.index.rename('date', inplace=True)

    #For each file-data, extract the basic stats, and add them to the df
    for fileData in filesData:
        firstMinuteAsleep = getFirstMinuteAsleep(fileData, minSleepIntervalRequired)
        stats = getSleepIntervalsStats(fileData, minSleepIntervalRequired)
        date = getSleepDate(fileData)
        intervalsStats.loc[date] = [firstMinuteAsleep, fileData.iloc[0]['datetime'].time(),
                                    fileData.iloc[-1]['datetime'].time(),
                                    stats.len.max(), stats.len.mean()]

    return intervalsStats

#Each row is a day with length of all its sleep intervals
def generateIntervalsStats(filesData, minSleepIntervalRequired=0):
    #Create new df to store the stats
    intervalsStats = pd.DataFrame(columns=np.arange(20))

    #For each file-data, extract the basic stats, and add them to the df
    for fileData in filesData:
        stats = getSleepIntervalsStats(fileData, minSleepIntervalRequired)
        date = getSleepDate(fileData)
        intervalsStats.loc[date] = stats.len

    return intervalsStats

#Each row is a complete intraday value array for the corresponding day
def generateIntradayStats(filesData, useTime=True):
    #Create new df to store the stats
    if useTime:
        minutes = pd.date_range('00:00', '23:59', freq='1min')
        intradayStats = pd.DataFrame(columns=[x.time().strftime("%H:%M")
                                               for x in minutes])
    else:
        intradayStats = pd.DataFrame(columns=np.arange(600))
    intradayStats.index.rename('date', inplace=True)

    #For each file-data, extract the basic stats, and add them to the df
    for fileData in filesData:
        data = fileData.copy(deep=True)
        date = getSleepDate(data)
        data['time'] = [x.strftime("%H:%M") for x in data['datetime'].dt.time]
        data.set_index(['time'], inplace=True)
        intradayStats.loc[date] = data['value']

    return intradayStats

#----------------------------#
#    SECOND LEVEL METHODS    #
#----------------------------#

def normalizedIntradayCountStats(intradayStats, limitCount=5):
    # For each minute, number of days for which we have a valid measure (record)
    notNullCount = intradayStats.count()
    # Ignore minutes where we have low level of records
    notNullCount[notNullCount < limitCount] = None
    # Count how many times each value appears for each minute
    valueCount = intradayStats.apply(pd.value_counts)
    # Normalize each minute by records count
    res = valueCount.div(notNullCount, axis=1)
    return res

def centerIntradayCountStats(intradayCountStats, timeToCenterOn="12:00"):
    columns = list(intradayCountStats.columns.values)
    k = columns.index(timeToCenterOn)
    columns = columns[k:] + columns[:k]
    return intradayCountStats[columns]

def getSleepIntervalsStats(data, minSleepIntervalRequired=0):
    sleepIntervals = getSleepIntervals(data, minSleepIntervalRequired)

    #Create new df to store the stats
    stats = pd.DataFrame(columns=['minute_start','minute_end','time_start','time_end','len'])

    #For each interval get start, end, len stats, and add it to the df
    for i, interval in enumerate(sleepIntervals):
        #Substract one because interval indicated the minute, and starts from 1
        timeStart = data.iloc[interval[0]-1]['datetime']
        timeEnd = data.iloc[interval[1]-1]['datetime']
        stats.loc[i] = [interval[0], interval[1], timeStart, timeEnd, (interval[1] - interval[0])+1]

    return stats

#Return the date associated with the provides sleep data.
#It should be the date associated with the first recorded minute, but if this is after
#midnight and before noon, we want to consider the previous day.
def getSleepDate(data):
    firstDatetime = data.ix[0]['datetime']
    if time(12,00) <= firstDatetime.time()<=time(23,59):
        return firstDatetime.date()
    else:
        return (firstDatetime - pd.DateOffset(1)).date()

def getFirstMinuteAsleep(data, minSleepIntervalRequired=0):
    firstAsleep = None

    values = data['value']
    indxs = values[values==SleepValue.sleeping.value].index
    for indx in indxs:
        if (len(values[indx:indx+minSleepIntervalRequired]) >= minSleepIntervalRequired and
                (values[indx:indx+minSleepIntervalRequired]==SleepValue.sleeping.value).all()):
            #Indexing starts from 0, minutes start from 1
            firstAsleep = indx+1
            break

    return firstAsleep

def getSleepIntervals(data, minSleepIntervalRequired=0):
    values = data['value']
    splits = np.append(np.where(np.diff(values)!=0)[0], [len(values)-1])+1

    prev = 0
    intervals = []
    for split in splits:
        if values[prev] == SleepValue.sleeping.value and (split-prev)>=minSleepIntervalRequired:
            intervals.append((prev+1, split))
        prev = split

    return intervals

def getMaxValues(basicStats, n, statName):
    return basicStats.nlargest(n, statName)[statName]


# -*- coding=utf-8 -*-
import cv2
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

img = cv2.imread('data/beauty.jpg')

# 通过cv2.cvtColor把图像从BGR转换到HSV
img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
# cv2.imshow('img_hsv', img_hsv)

# H空间中，绿色比黄色的值高一点，所以给每个像素+15，黄色的树叶就会变绿
turn_green_hsv = img_hsv.copy()
turn_green_hsv[:, :, 0] = (turn_green_hsv[:, :, 0] + 15) % 180
turn_green_img = cv2.cvtColor(turn_green_hsv, cv2.COLOR_HSV2BGR)
cv2.imwrite('output/beauty_green.jpg', turn_green_img)
# cv2.imshow('img_hsv_turn_green', turn_green_img)

# 减小饱和度会让图像损失鲜艳，变得更灰
colorless_hsv = img_hsv.copy()
colorless_hsv[:, :, 1] = 0.5 * colorless_hsv[:, :, 1]
colorless_img = cv2.cvtColor(colorless_hsv, cv2.COLOR_HSV2BGR)
cv2.imwrite('output/beauty_colorless.jpg', colorless_img)
# cv2.imshow('img_hsv_turn_gray', colorless_img)

# 减小明度为原来一半
darker_hsv = img_hsv.copy()
darker_hsv[:, :, 2] = 0.5 * darker_hsv[:, :, 2]
darker_img = cv2.cvtColor(darker_hsv, cv2.COLOR_HSV2BGR)
cv2.imwrite('output/beauty_darker.jpg', darker_img)
# cv2.imshow('darker_img', darker_img)

########### Gamma ############
# 分通道计算每个通道的直方图
hist_b = cv2.calcHist([img], [0], None, [256], [0, 256])
hist_g = cv2.calcHist([img], [1], None, [256], [0, 256])
hist_r = cv2.calcHist([img], [2], None, [256], [0, 256])

# 定义Gamma矫正的函数
def gamma_trans(img, gamma):
	# 具体做法是先归一化到1，然后gamma作为指数值求出新的像素值再还原
	gamma_table = [np.power(x/255.0, gamma)*255.0 for x in range(256)]
	gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)

	# 实现这个映射用的是OpenCV的查表函数
	return cv2.LUT(img, gamma_table)

# 执行Gamma矫正，小于1的值让暗部细节大量提升，同时亮部细节少量提升
img_corrected = gamma_trans(img, 0.5)
cv2.imwrite('output/beauty_gamma_corrected.jpg', img_corrected)
cv2.imshow('gamma_corrected.jpg', img_corrected)

# 分通道计算Gamma矫正后的直方图
hist_b_corrected = cv2.calcHist([img_corrected], [0], None, [256], [0, 256])
hist_g_corrected = cv2.calcHist([img_corrected], [1], None, [256], [0, 256])
hist_r_corrected = cv2.calcHist([img_corrected], [2], None, [256], [0, 256])

# 将直方图进行可视化
fig = plt.figure()

pix_hists = [
	[hist_b, hist_g, hist_r],
	[hist_b_corrected, hist_g_corrected, hist_r_corrected]
]

pix_vals = range(256)
for sub_plt, pix_hist in zip([121, 122], pix_hists):
	ax = fig.add_subplot(sub_plt, projection = '3d')
	for c, z, channel_hist in zip(['b', 'g', 'r'], [20, 10, 0], pix_hist):
		cs = [c] * 256
		ax.bar(pix_vals, channel_hist, zs = z, zdir = 'y', color = cs, alpha = 0.618, edgecolor = 'none', lw = 0)

	ax.set_xlabel('Pixel Values')
	ax.set_xlim([0, 256])
	ax.set_ylabel('Counts')
	ax.set_zlabel('Channels')

plt.show()

cv2.waitKey()


import numpy as np
import qutip
import scipy.sparse.coo as coo
import time
import sys
import qc
import tabulate

FILE_QC0 = 'table_qc0'
FILE_QCF = 'table_qcf'
PATH_TEX = '../doc/%s.tex'
PATH_CSV = '../doc/csv/%s.csv'

tabulate.LATEX_ESCAPE_RULES = {}

class QMemProfiler:
	def __init__(self):
		table_sizes = []
		self.profile_mem()

	def mem_qc0(self, N):
		qs = qc.QSimon({"bits":N})
		sizes = qs.sz
		#print('Memory QC0: {}'.format(sizes))
		return qs, sizes

	def mem_qcf(self, qs, config):
		st = qs.run(config)
		sizes = qs.sz
		#print('Memory QCf: {}'.format(sizes))
		return st, sizes

	def profile_mem(self):
		Nmin = 2
		Nmax = 10

		table_qc0 = []
		table_qcf = []
		headers_qc0 = ['n', 'N', 'log2_H', 'log2_phi1', 'log2_all_qc0', 'log2_approx_qc0']
		headers_qcf = ['n', 'N', 'log2_H', 'log2_U', 'log2_qstate', 'log2_all_qcf' ,'log2_approx_qcf']
		tex_headers_qc0 = [
			'$n$',
			'$N$',
			'$\log_2 S(H)$',
			'$\log_2 S(\ket{\phi_1})$',
			'$\log_2 S_T$',
			'$\log_2 S_T\'$'
		]
		tex_headers_qcf = [
			'$n$',
			'$N$',
			'$\log_2 S(H)$',
			'$\log_2 S(U)$',
			'$\log_2 S(\ket{\phi_3})$',
			'$\log_2 S_T$'
			,'$\log_2 S_T\'$'
		]

		for N in range(Nmin, Nmax+1):
			entry = {}
			row_qc0 = []
			row_qcf = []

			print("N = {}".format(N))
			qs, size_qc0 = self.mem_qc0(N)
			config = qs.random_config(N)
			st, size_qcf = self.mem_qcf(qs, config)
			# Mostrar H y phi1 en qc0
			# U y qstate en qcf
			entry['N']				= 2*N
			entry['n']				= N
			entry['H']				= size_qc0['H']
			entry['phi1']			= size_qc0['phi1']
			entry['all_qc0']		= size_qc0['all']
			entry['log2_H']			= np.log2(size_qc0['H'])
			entry['log2_phi1']		= np.log2(size_qc0['phi1'])
			entry['log2_all_qc0']	= np.log2(size_qc0['all'])
			entry['log2_approx_qc0'] = np.log2(2**(2*N+3) + 2**(2*N+2))
			# El estado phi1 puede ser machacado?
			entry['U']				= size_qcf['U']
			entry['qstate']			= size_qcf['qstate']
			entry['all_qcf']		= size_qcf['U'] + size_qcf['H'] + size_qcf['qstate']
			entry['log2_U']			= np.log2(size_qcf['U'])
			entry['log2_qstate']	= np.log2(size_qcf['qstate'])
			entry['log2_all_qcf'] = np.log2(
				size_qcf['U'] + size_qcf['H'] + size_qcf['qstate'])
			entry['log2_approx_qcf'] = np.log2(
#				2**(2*N+3) + 2**(2*N+2) +	#U
#				2**(2*N+3) +				#H
#				2**(2*N+2) + 2**(2*N+1)		#qstate
				2**(2*N+3) + 2**(2*N+2) +
				2**(2*N+3) +
				2**(2*N+2) + 2**(2*N+1)
			)

			for h in headers_qc0:
				row_qc0.append(entry[h])
			for h in headers_qcf:
				row_qcf.append(entry[h])

			table_qc0.append(row_qc0)
			table_qcf.append(row_qcf)

		self.save_tex(table_qc0, tex_headers_qc0, FILE_QC0)
		self.save_tex(table_qcf, tex_headers_qcf, FILE_QCF)

		self.save_csv(table_qc0, headers_qc0, FILE_QC0)
		self.save_csv(table_qcf, headers_qcf, FILE_QCF)

	def save_csv(self, t, h, fn):
		str_header = ','.join(h)
		arrt = np.asarray(t)

		np.savetxt(PATH_CSV % fn, arrt, header=str_header, delimiter=",",
			comments='')

	def save_tex(self, t, h, fn):
		float_fmt = ".2f"
		table_fmt = 'latex_booktabs'
		str_table = tabulate.tabulate(t, h, floatfmt=float_fmt,
			tablefmt=table_fmt)

		self.save_file(str_table, PATH_TEX % fn)

	def save_file(self, s, fn):
		f = open(fn, 'w')
		f.write(s)
		f.close()

QMemProfiler()


import numpy as np
import math

class Verifier():

    def __init__(self, random_example_generator, params_generator,  epsilon=0.05, delta=0.05):
        assert((epsilon <= 1) & (delta <= 1))
        self.epsilon = epsilon
        self.delta = delta
        self._log_delta = np.log(delta)
        self._log_one_minus_epsilon = np.log(1-epsilon)
        self._num_equivalence_asked = 0
        self._last_counterexample_positive = False
        self._num_counterexamples_in_EQ = None
        self._number_of_samples = None
        self.number_of_examples_checked = 0
        self._get_random_example = random_example_generator
        self._params_generator = params_generator
        self.filtered_by_query = 0

    def equivalence_check(self, blackbox, learner, query, verbose = False):

        self._num_equivalence_asked += 1
        self._number_of_samples = int(
            math.ceil((self._num_equivalence_asked*0.693147-self._log_delta)/self.epsilon))

        _found_a_counterexample = None

        # generator
        generator = self._get_random_example(X=self._params_generator[0], feature_type=self._params_generator[1], size=self._number_of_samples)

        # typically it is ideal to get a random sample and then process it,
        # but SyGuS works better (its prediction function) when given a set of examples for classification
        examples = []
        for example in generator:

            # when the example is filtered out by the query, it is not considered for further analysis
            if(not query.classify_example(example)):
                self.filtered_by_query += 1
                continue

            examples.append(example)

        # get classification from learner and blackbox
        if(len(examples) > 0):
            learner_verdicts = learner.classify_examples(examples)
            blackbox_verdicts = blackbox.classify_examples(examples)

        for i in range(len(examples)):

            if(learner_verdicts != None):
                learner_verdict = learner_verdicts[i]
            else:
                learner_verdict = learner_verdicts

            self.number_of_examples_checked += 1

            if(learner_verdict == None or (learner_verdict != blackbox_verdicts[i])):

                _found_a_counterexample = True

                # toggle between positive and negative counterexamples
                if(blackbox_verdicts[i] == self._last_counterexample_positive):
                    self._last_counterexample_positive =  not self._last_counterexample_positive
                else:
                    # store one counterexample in case toggle does not work
                    _found_a_counterexample = (examples[i], blackbox_verdicts[i])
                    continue

                if(verbose):
                    print("-found", blackbox_verdicts[i] == True, "counterexample")

                # print(learner_verdict, blackbox_verdict, query_verdict, 1 - learner_verdict)
                return examples[i], blackbox_verdicts[i]

        if(_found_a_counterexample is not None):
            example, label = _found_a_counterexample
            if(verbose):
                print("-could not find", self._last_counterexample_positive, "counterexample. Only found", label)
            return example, label

        return None, None


# adapted from Distill-VQ
import os.path as osp
import numpy as np
from collections import defaultdict
from tqdm import tqdm
import faiss
import pickle
import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist
from scipy.optimize import linear_sum_assignment

class ProductQuantization(nn.Module):
    def __init__(
        self,
        pq_type: str = 'pq',
        subvector_num: int = 32,
        subvector_bits: int = 8,
        dist_mode: str = 'ip',
        emb_size: int = 768,
        pq_init_method: str = 'faiss',
        pq_update_method: str = 'grad',
        tie_nci_pq_centroid: int = 0,
        lm_head: nn.Parameter = None,
        centroid_update_loss: str = 'none',
        rq_topk_score: str = 'prod',
    ):
        super().__init__()
        assert pq_type in ('pq', 'opq', 'rq')
        assert dist_mode in ('ip', 'l2', 'iptol2')
        # dist_mode 'iptol2' means using l2 to represent ip
        # https://gist.github.com/mdouze/e4bdb404dbd976c83fe447e529e5c9dc
        # query (q1,...,qd) to (q1,...,qd,0)
        # document (p1,...,pd) to (p1,...,pd,sqrt(PHI-l2norm^2))
        self.pq_type = pq_type
        self.subvector_num = subvector_num
        self.subvector_bits = subvector_bits
        self.subvector_cents = 2 ** subvector_bits
        self.dist_mode = dist_mode
        self.emb_size = emb_size
        self.pq_init_method = pq_init_method
        self.pq_update_method = pq_update_method
        self.tie_nci_pq_centroid = tie_nci_pq_centroid
        self.centroid_update_loss = centroid_update_loss
        self.rq_topk_score = rq_topk_score
        self.get_preds = False

        if self.pq_type == 'rq':
            last_dim = emb_size
        else:
            last_dim = emb_size // subvector_num
        self.last_dim = last_dim

        if self.tie_nci_pq_centroid:
            assert self.dist_mode != 'iptol2'
            assert self.pq_type != 'opq'
            self.lm_head = lm_head
            if self.pq_type == 'pq':
                self.weight_layers = nn.Sequential(
                    nn.Linear(emb_size, emb_size), nn.ReLU(), nn.Linear(emb_size, last_dim))
        else:
            codebook_last_dim = last_dim
            if self.dist_mode == 'iptol2':
                codebook_last_dim += 1
            self.codebook = nn.Parameter(torch.empty(
                subvector_num, self.subvector_cents, codebook_last_dim), requires_grad=(self.pq_update_method == 'grad')).type(torch.FloatTensor)
            if self.pq_type == 'opq':
                self.rotate = nn.Parameter(torch.empty(
                    emb_size, emb_size), requires_grad=False).type(torch.FloatTensor)
            if self.pq_update_method == 'ema':
                # adapted from rq-vae-transformer
                self.decay = 0.99
                self.eps = 1e-5
                self.restart_unused_codes = True
                self.register_buffer('cluster_size_ema', torch.zeros(
                    subvector_num, self.subvector_cents))
                self.register_buffer(
                    'embed_ema', self.codebook.detach().clone())

    def augment_xb(self, xb, phi=None):
        norms = np.sum((xb ** 2), axis=-1)
        if phi is None:
            phi = np.max(norms)
        extracol = np.sqrt(phi - norms)
        return np.hstack((xb, extracol[..., np.newaxis]))

    def augment_xq(self, xq):
        if isinstance(xq, torch.Tensor):
            xq = torch.cat((xq, xq.new_zeros(*xq.shape[:-1], 1)), dim=-1)
        else:
            xq = np.concatenate(
                (xq, np.zeros((*xq.shape[:-1], 1), dtype=xq.dtype)), axis=-1)
        return xq

    def wrapped_augment_xb(self, xb, index=None):
        if self.dist_mode != 'iptol2':
            return xb
        emb_size = self.last_dim
        reshaped = False
        if xb.shape[-1] not in (emb_size, emb_size + 1):
            ori_shape = xb.shape
            xb = xb.reshape(-1, emb_size)
            reshaped = True
        if xb.shape[-1] == emb_size:
            xb = self.augment_xb(xb)
            extracol = xb[..., -1]
            if reshaped:
                xb = xb.reshape(*ori_shape[:-1], -1)
        else:
            extracol = xb[..., -1]
        assert xb.shape[-1] == emb_size + 1
        if index is None:
            self.extracol[:] = torch.tensor(extracol, dtype=torch.float32)
        else:
            self.extracol[:, index] = torch.tensor(
                extracol, dtype=torch.float32)
        return xb

    def rq_minus_centroids(self, embeddings, centroids):
        return embeddings - centroids[..., :self.last_dim]

    def compute_scores(self, a, b):
        if self.dist_mode == 'ip':
            result = a * b
        else:
            if self.dist_mode == 'iptol2' and a.shape[-1] != b.shape[-1]:
                a = self.augment_xq(a)
            result = - ((a - b) ** 2)
        return torch.sum(result, dim=-1)

    def get_codebook(self):
        if self.tie_nci_pq_centroid:
            codebook = self.lm_head[2:2+self.subvector_num*self.subvector_cents].reshape(
                self.subvector_num, self.subvector_cents, self.emb_size)
            if self.pq_type == 'pq':
                codebook = self.weight_layers(codebook)
            return codebook
        else:
            return self.codebook

    def codebook_from_index(self, index, index_file=None):
        if index is None:
            index = faiss.read_index(index_file)
        if self.pq_type == 'opq':
            assert isinstance(index, faiss.IndexPreTransform)
            vt = faiss.downcast_VectorTransform(index.chain.at(0))
            assert isinstance(vt, faiss.LinearTransform)
            rotate = faiss.vector_to_array(vt.A).reshape(vt.d_out, vt.d_in)
            index = faiss.downcast_index(index.index)
        else:
            rotate = None
            index = index

        if self.pq_type == 'rq':
            rq = index.rq
            centroid_embeds = faiss.vector_to_array(rq.codebooks)
            codebook = centroid_embeds.reshape(
                rq.M, rq.total_codebook_size // rq.M, -1)
            assert self.last_dim == rq.d
        else:
            pq_index = index
            centroid_embeds = faiss.vector_to_array(pq_index.pq.centroids)
            codebook = centroid_embeds.reshape(
                pq_index.pq.M, pq_index.pq.ksub, -1)
            assert self.last_dim == pq_index.pq.dsub
        assert self.codebook.shape[-1] == codebook.shape[-1], 'Codebook shape not aligned! Considering whether iptol2.'

        with torch.no_grad():
            self.codebook.copy_(torch.tensor(codebook))
            if self.pq_type == 'opq':
                self.rotate.copy_(torch.tensor(rotate))

    def build_faiss_index(self, doc_embeddings, save_path=None):
        assert self.dist_mode in ('ip', 'l2', 'iptol2')
        index_metric = faiss.METRIC_INNER_PRODUCT if self.dist_mode == 'ip' else faiss.METRIC_L2

        if self.pq_type == 'pq':
            components = f"PQ{self.subvector_num}x{self.subvector_bits}"
        elif self.pq_type == 'opq':
            components = f"OPQ{self.subvector_num},PQ{self.subvector_num}x{self.subvector_bits}"
        elif self.pq_type == 'rq':
            components = f"RQ{self.subvector_num}x{self.subvector_bits}"
        if self.dist_mode == 'iptol2':
            assert self.pq_type != 'rq', 'Not support iptol2 rq using faiss index.'
            doc_embeddings = self.wrapped_augment_xb(doc_embeddings)
        emb_size = doc_embeddings.shape[-1]

        index = faiss.index_factory(
            emb_size, components, index_metric)

        index.train(doc_embeddings)
        index.add(doc_embeddings)

        if save_path is not None:
            faiss.write_index(index, save_path)
        return index

    @torch.no_grad()
    def get_document_cluster_simple(self, return_mapping: bool = False):
        assert self.get_preds
        cluster = defaultdict(list)
        mapping = {}
        for i, p in enumerate(self.last_preds):
            label = tuple(p.tolist())
            mapping[i] = label
            cluster[label].append(i)
        del self.last_preds
        self.get_preds = False
        if return_mapping:
            return dict(cluster), mapping
        else:
            return dict(cluster)

    @torch.no_grad()
    def get_document_cluster(self, doc_embeddings: np.ndarray, rank: int, nrank: int, batch_size: int = 1024, return_mapping: bool = False):
        num_docs = doc_embeddings.shape[0]
        num_docs_per_worker = num_docs // nrank
        start = num_docs_per_worker * rank
        if rank + 1 == nrank:
            ending = num_docs
        else:
            ending = start + num_docs_per_worker
        part_num_docs = ending - start
        cluster = torch.empty(
            (part_num_docs, self.subvector_num), dtype=torch.int32)
        if self.pq_type == 'rq':
            func = self.get_rq_document_cluster
        else:
            func = self.get_pq_document_cluster
        func(doc_embeddings, cluster, start,
             ending, rank, batch_size)
        if return_mapping:
            new_mapping = {}
        doc_cluster = defaultdict(list)
        for k, v in enumerate(cluster):
            tuplev = tuple(v.tolist())
            key = k + start
            doc_cluster[tuplev].append(key)
            if return_mapping:
                new_mapping[key] = tuplev
        print('Number of document clusters:', len(doc_cluster))
        if return_mapping:
            return dict(doc_cluster), new_mapping
        else:
            return dict(doc_cluster)

    def get_pq_document_cluster(self, doc_embeddings: np.ndarray, cluster: torch.Tensor, start: int, ending: int, rank: int, batch_size: int = 1024):
        emb_size = doc_embeddings.shape[1]
        part_emb_size = emb_size // self.subvector_num
        codebook = self.get_codebook().cpu()
        with torch.no_grad():
            for i in tqdm(range(start, ending, batch_size), desc=f'Document Cluster {rank}'):
                batch_start = i
                batch_ending = min(i+batch_size, ending)
                cur_embeddings = torch.tensor(
                    doc_embeddings[batch_start:batch_ending], device='cpu')
                if self.pq_type == 'opq':
                    cur_embeddings = torch.matmul(
                        cur_embeddings, self.rotate.cpu().T)
                cur_codebook = codebook.unsqueeze(
                    1).expand(-1, cur_embeddings.size(0), -1, -1)
                if self.dist_mode == 'iptol2':
                    cur_extracol = self.extracol[batch_start:batch_ending]
                for j in range(self.subvector_num):
                    emb_start = j*part_emb_size
                    emb_ending = emb_start + part_emb_size
                    part_doc_embeddings = cur_embeddings[:,
                                                         emb_start:emb_ending]
                    if self.dist_mode == 'iptol2':
                        part_doc_embeddings = torch.cat(
                            (part_doc_embeddings, cur_extracol[:, j].unsqueeze(-1)), -1)
                    part_doc_embeddings = part_doc_embeddings.unsqueeze(-2)
                    part_codebook = cur_codebook[j]
                    distance = self.compute_scores(
                        part_doc_embeddings, part_codebook)
                    index = distance.max(dim=-1, keepdim=True)[1].squeeze(-1)
                    cluster[batch_start - start: batch_ending - start, j] = index

    def get_rq_document_cluster(self, doc_embeddings: np.ndarray, cluster: torch.Tensor, start: int, ending: int, rank: int, batch_size: int = 1024):
        part_num_docs = ending - start
        temp_docemb = torch.tensor(doc_embeddings[start:ending], device='cpu')
        codebook = self.get_codebook()
        with torch.no_grad():
            for j in range(self.subvector_num):
                part_codebook = codebook[j].cpu()
                if self.dist_mode == 'iptol2':
                    cur_extracol = self.extracol[:, j].unsqueeze(-1)
                for i in tqdm(range(0, part_num_docs, batch_size), desc=f'Document Cluster {rank}-{j}'):
                    batch_start = i
                    batch_ending = min(i+batch_size, part_num_docs)
                    cur_embeddings = temp_docemb[batch_start:batch_ending]
                    if self.dist_mode == 'iptol2':
                        cur_embeddings = torch.cat(
                            (cur_embeddings, cur_extracol[batch_start:batch_ending]), -1)
                    cur_embeddings = cur_embeddings.unsqueeze(-2)
                    cur_codebook = part_codebook.unsqueeze(
                        0).expand(cur_embeddings.size(0), -1, -1)
                    distance = self.compute_scores(
                        cur_embeddings, cur_codebook)
                    index = distance.max(dim=-1, keepdim=True)[1].squeeze(-1)
                    cluster[batch_start:batch_ending, j] = index
                    temp_docemb[batch_start:batch_ending] = self.rq_minus_centroids(
                        temp_docemb[batch_start:batch_ending], part_codebook[index])

    def forward(self, vecs, return_loss=True):
        # TODO: check results and determine which is better
        # before commit we generate label AFTER gumbel softmax
        # after commit we generate label BEFORE gumbel softmax
        if self.pq_type == 'rq':
            proba, index, loss = self.forward_rq(vecs, return_loss)
        else:
            proba, index, loss = self.forward_pq(vecs, return_loss)
        if self.training and self.pq_update_method == 'ema':
            self.ema_update(vecs, index)
        return proba, index, loss

    def forward_pq(self, vecs, return_loss=True):
        if self.pq_type == 'opq':
            vecs = torch.matmul(vecs, self.rotate.T)
        vecs = vecs.view(vecs.size(0), self.subvector_num, -1)
        codebook = self.get_codebook().unsqueeze(0).expand(vecs.size(0), -1, -1, -1)
        proba = self.compute_scores(vecs.unsqueeze(-2), codebook)
        index = proba.max(dim=-1)[1]
        if return_loss:
            if self.centroid_update_loss == 'reconstruct':
                reconstruct_emb = self.get_reconstruct_vector(index, codebook).view(
                    index.shape[0], self.subvector_num, self.last_dim)
                loss = ((vecs - reconstruct_emb) ** 2).mean()
            else:
                loss = None
        else:
            loss = None
        return proba, index, loss

    def forward_rq(self, vecs, return_loss=True):
        allproba = []
        index = []
        codebook = self.get_codebook()
        use_reconstruct_loss = self.centroid_update_loss == 'reconstruct'
        if use_reconstruct_loss:
            errors = []
        for i in range(self.subvector_num):
            cur_codebook = codebook[i:i+1].expand(vecs.size(0), -1, -1)
            proba = self.compute_scores(vecs.unsqueeze(-2), cur_codebook)
            part_index = proba.max(dim=-1)[1]
            allproba.append(proba)
            index.append(part_index)
            if self.dist_mode == 'iptol2':
                cur_centroid = codebook[i][part_index][..., :-1]
            else:
                cur_centroid = codebook[i][part_index]
            if use_reconstruct_loss:
                cur_error = vecs.detach() - cur_centroid
                errors.append(cur_error)
            if i != self.subvector_num - 1:
                vecs -= cur_centroid.detach()
        proba = torch.stack(allproba, dim=1)
        index = torch.stack(index, dim=1)
        if return_loss:
            if use_reconstruct_loss:
                loss = (torch.stack(errors) ** 2).mean()
                # loss = (vecs ** 2).mean()
            else:
                loss = None
        else:
            loss = None
        return proba, index, loss

    @torch.no_grad()
    def ema_update(self, vectors, idxs):
        assert self.dist_mode != 'iptol2'
        embed_dim = self.last_dim
        if self.pq_type == 'rq':
            vectors = vectors.unsqueeze(
                1).expand(-1, self.subvector_num, -1).reshape(-1, embed_dim)
        else:
            if self.pq_type == 'opq':
                vectors = torch.matmul(vectors, self.rotate.T)
            vectors = vectors.reshape(-1, embed_dim)
        idxs = idxs.reshape(-1)
        n_vectors = vectors.shape[0]
        one_hot_idxs = vectors.new_zeros(n_vectors, self.subvector_cents)
        one_hot_idxs.scatter_(dim=1,
                              index=idxs.unsqueeze(1),
                              src=vectors.new_ones(n_vectors, 1)
                              )
        one_hot_idxs = one_hot_idxs.reshape(-1,
                                            self.subvector_num, self.subvector_cents)
        cluster_size = one_hot_idxs.sum(dim=0)
        vectors_sum_per_cluster = torch.bmm(one_hot_idxs.permute(
            1, 2, 0), vectors.reshape(-1, self.subvector_num, embed_dim).permute(1, 0, 2))

        if dist.is_initialized():
            dist.all_reduce(vectors_sum_per_cluster, op=dist.ReduceOp.SUM)
            dist.all_reduce(cluster_size, op=dist.ReduceOp.SUM)

        self.cluster_size_ema.mul_(self.decay).add_(
            cluster_size, alpha=1 - self.decay)
        self.embed_ema.mul_(self.decay).add_(
            vectors_sum_per_cluster, alpha=1 - self.decay)

        if self.restart_unused_codes:
            temp = vectors.reshape(-1, self.subvector_num, embed_dim)
            B = temp.shape[0]
            if B < self.subvector_cents:
                n_repeats = (self.subvector_cents + B - 1) // B
                std = temp.new_ones(embed_dim) * 0.01 / np.sqrt(embed_dim)
                temp = temp.repeat(n_repeats, 1, 1)
                temp = temp + torch.rand_like(temp) * std

            _vectors_random = torch.stack([temp[torch.randperm(
                temp.shape[0], device=temp.device), i][:self.subvector_cents] for i in range(self.subvector_num)])

            if dist.is_initialized():
                dist.broadcast(_vectors_random, 0)

            usage = (self.cluster_size_ema.unsqueeze(-1) >= 1).float()
            self.embed_ema.mul_(usage).add_(_vectors_random * (1-usage))
            usage = usage.squeeze(-1)
            self.cluster_size_ema.mul_(usage)
            self.cluster_size_ema.add_(torch.ones_like(
                self.cluster_size_ema) * (1-usage))

        # update embedding
        n = self.cluster_size_ema.sum(dim=1, keepdim=True)
        normalized_cluster_size = (
            n * (self.cluster_size_ema + self.eps) /
            (n + self.subvector_cents * self.eps)
        )
        self.codebook.data[:] = self.embed_ema / \
            normalized_cluster_size.unsqueeze(-1)

    def fix(self):
        self.codebook.requires_grad_(False)
        if self.pq_type == 'opq':
            self.rotate.requires_grad_(False)

    @torch.no_grad()
    def initialize(self, index_file, doc_emb, rank, seed, pq_cluster_path, encode_batch_size):
        if self.dist_mode == 'iptol2':
            self.extracol = torch.empty(
                (doc_emb.shape[0], self.subvector_num), dtype=torch.float32)
        if self.pq_init_method != 'none':
            use_file = index_file is not None and osp.isfile(index_file)
            if self.pq_init_method.endswith('kmeans') and not use_file:
                self.get_preds = True
            if rank == 0:
                if self.pq_init_method == 'faiss':
                    if use_file:
                        print("Intializing codebook with faiss index file...")
                        index = faiss.read_index(index_file)
                    else:
                        print("Intializing codebook by building faiss index...")
                        index = self.build_faiss_index(
                            doc_emb, save_path=index_file)
                    self.codebook_from_index(index)
                elif self.pq_init_method.endswith('kmeans'):
                    if use_file:
                        print("Intializing codebook with torch file...")
                        tensor = torch.load(index_file, map_location='cpu')
                        self.codebook.copy_(tensor)
                        del tensor
                    else:
                        print("Intializing codebook by kmeans clustering...")
                        self.unsupervised_update_codebook_manually(
                            doc_emb, seed, self.pq_init_method)
                        if index_file is not None:
                            torch.save(self.codebook, index_file)
                elif self.pq_init_method == 'avg':
                    if use_file:
                        tensor = torch.load(index_file, map_location='cpu')
                        self.codebook.copy_(tensor)
                        del tensor
                    else:
                        torch.nn.init.normal_(
                            self.codebook.data, mean=0.0, std=0.01)
                    print(
                        f"Intializing codebook using average document embedding after {use_file} use file...")
                    self.init_pq_using_document_cluster(
                        doc_emb, pq_cluster_path, encode_batch_size)
            if dist.is_initialized():
                dist.broadcast(self.codebook.data, 0)
                if self.pq_type == 'opq':
                    dist.broadcast(self.rotate.data, 0)

    @torch.no_grad()
    def init_pq_using_document_cluster(self, doc_emb, cluster, batch_size):
        assert self.dist_mode in ('l2', 'iptol2')
        assert self.pq_type in ('pq', 'rq')
        with open(cluster, 'rb') as fr:
            cluster = pickle.load(fr)
        if self.dist_mode == 'iptol2':
            doc_emb = self.wrapped_augment_xb(doc_emb)
        doc_emb = np.array(doc_emb)

        macro_cluster = [defaultdict(list) for _ in range(self.subvector_num)]
        for k, v in cluster.items():
            for i in range(self.subvector_num):
                macro_cluster[i][k[i]].extend(v)
        macro_cluster = [dict(mc) for mc in macro_cluster]
        if self.pq_type == 'rq':
            part_doc_emb = doc_emb
        for i in range(self.subvector_num):
            cur_cluster = macro_cluster[i]
            if self.pq_type == 'pq':
                dstart = i * self.last_dim
                dending = dstart + self.last_dim
                part_doc_emb = doc_emb[:, dstart:dending]
            cur_codebook = self.codebook[i]
            for k, v in tqdm(cur_cluster.items(), desc=f'Averaging Embedding'):
                accum = np.zeros(self.last_dim)
                ndocs = len(v)
                for start in range(0, ndocs, batch_size):
                    docs = v[start:start+batch_size]
                    cur_emb = part_doc_emb[docs]
                    accum += (np.sum(cur_emb, axis=0) / ndocs)
                cur_codebook[k][:] = torch.tensor(accum)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    for start in range(0, ndocs, batch_size):
                        docs = v[start:start+batch_size]
                        part_doc_emb[docs] -= accum
        # torch.save(self.codebook, 'newcodebook.pt')

    @torch.no_grad()
    def unsupervised_update_codebook(self, doc_emb, rank, seed, align=False):
        if self.pq_update_method.endswith('kmeans'):
            self.get_preds = True
        if rank == 0:
            if align:
                ori_codebook = self.codebook.data.clone()
            if self.pq_update_method == 'faiss':
                self.unsupervised_update_codebook_faiss(doc_emb, seed)
            elif self.pq_update_method.endswith('kmeans'):
                self.unsupervised_update_codebook_manually(
                    doc_emb, seed, self.pq_update_method)
            if align:
                self.align_codebook(ori_codebook)
        dist.broadcast(self.codebook.data, 0)
        if self.pq_type == 'opq':
            dist.broadcast(self.rotate.data, 0)

    @torch.no_grad()
    def unsupervised_update_codebook_faiss(self, doc_emb, seed):
        print("Updating codebook using Faiss...")
        index = self.build_faiss_index(doc_emb)
        self.codebook_from_index(index)

    @torch.no_grad()
    def unsupervised_update_codebook_manually(self, doc_emb, seed, kmeans_method):
        print("Updating codebook using KMeans...")
        assert self.pq_type != 'opq'
        emb_size = doc_emb.shape[1]
        k = 2 ** self.subvector_bits
        if kmeans_method == 'kmeans':
            from sklearn.cluster import KMeans, MiniBatchKMeans
            if doc_emb.shape[0] >= 1e3:
                kmeans = MiniBatchKMeans(n_clusters=k, max_iter=300, n_init=100, init='k-means++', random_state=seed,
                                         batch_size=1000, reassignment_ratio=0.01, max_no_improvement=20, tol=1e-7)
            else:
                kmeans = KMeans(n_clusters=k, max_iter=300, n_init=100,
                                init='k-means++', random_state=seed, tol=1e-7)
        elif kmeans_method == 'balancekmeans':
            raise NotImplementedError
        else:
            assert False

        preds = []
        if self.pq_type == 'pq':
            part_emb_size = emb_size // self.subvector_num
            codebook = []
            for i in tqdm(range(self.subvector_num), 'KMeans Clustering'):
                start = i * part_emb_size
                ending = start + part_emb_size
                cur_docemb = doc_emb[:, start:ending]
                cur_docemb = self.wrapped_augment_xb(cur_docemb, i)
                pred = kmeans.fit_predict(cur_docemb)
                codebook.append(kmeans.cluster_centers_)
                preds.append(pred)
            codebook = np.stack(codebook)
        elif self.pq_type == 'rq':
            codebook = []
            temp_docemb = np.array(doc_emb)
            for i in tqdm(range(self.subvector_num), 'KMeans Clustering'):
                temp_docemb = self.wrapped_augment_xb(temp_docemb, i)
                pred = kmeans.fit_predict(temp_docemb)
                centers = kmeans.cluster_centers_
                codebook.append(centers)
                preds.append(pred)
                if i != self.subvector_num - 1:
                    temp_docemb = self.rq_minus_centroids(
                        temp_docemb[..., :self.last_dim], centers[pred])
            codebook = np.stack(codebook)
        preds = np.stack(preds, axis=-1)
        self.last_preds = preds
        with torch.no_grad():
            self.codebook.copy_(torch.tensor(codebook))

    @torch.no_grad()
    def align_codebook(self, ori_codebook):
        # use sort; only need to sort 256^2 numbers using 8 bits
        new_codebook = self.codebook.new(*self.codebook.shape)
        for ori, cur, new in zip(ori_codebook, self.codebook, new_codebook):
            scores = self.compute_scores(ori.unsqueeze(0), cur.unsqueeze(1))
            scores = scores.cpu().numpy()
            assign = linear_sum_assignment(scores, maximize=True)
            for cid, oid in zip(*assign):
                new[oid] = cur[cid]
        with torch.no_grad():
            self.codebook.copy_(new_codebook)

    @torch.no_grad()
    def beam_search(
        self,
        doc_emb: torch.Tensor,
        num_return_sequences,
        num_beams=None,
        do_sample=False,
        return_proba=False,
    ):
        if num_beams is None:
            num_beams = num_return_sequences

        # (nsub, vocab, dim)
        codebook = self.get_codebook()
        if doc_emb.device != codebook.device:
            doc_emb = doc_emb.to(codebook.device)
        if self.pq_type == 'opq':
            doc_emb = torch.matmul(doc_emb, self.rotate.T)
        batch_size = doc_emb.size(0)
        # (bs, beam)
        beam_scores = doc_emb.new_ones(batch_size, 1)
        if self.pq_type == 'rq':
            # (bs, beam, dim)
            temp_embed = doc_emb.unsqueeze(1).clone()
        # (bs, beam, nsub)
        temp_index = torch.zeros(
            (batch_size, 1, 1), device=doc_emb.device, dtype=torch.int32)
        # (beam * vocab)
        beam_indices = torch.div(torch.arange(
            num_beams * self.subvector_cents, device=doc_emb.device), self.subvector_cents, rounding_mode='floor')
        code_indices = torch.arange(
            self.subvector_cents, device=doc_emb.device).repeat(num_beams)

        for i in range(self.subvector_num):
            if self.pq_type == 'rq':
                # (bs, 1, vocab, dim)
                cur_codebook = codebook[i:i +
                                        1].expand(batch_size, -1, -1).unsqueeze(1)
                cur_embed = temp_embed.unsqueeze(-2)
            else:
                # (bs, vocab, dim)
                cur_codebook = codebook[i].unsqueeze(
                    0).expand(batch_size, -1, -1)
                start = i * self.last_dim
                ending = start + self.last_dim
                # (bs, dim)
                cur_embed = doc_emb[:, start:ending].unsqueeze(1)
            # (bs, beam, vocab)
            proba = self.compute_scores(cur_embed, cur_codebook)
            proba = F.softmax(proba, dim=-1)
            if self.pq_type == 'rq':
                if self.rq_topk_score == 'prod':
                    proba = beam_scores.unsqueeze(-1) * proba
                else:
                    proba = proba
            else:
                proba = beam_scores.unsqueeze(-1) * proba.unsqueeze(1)
            # (bs, beam * vocab)
            proba = proba.view(batch_size, -1)
            prev_nbeam = beam_scores.size(1)
            if prev_nbeam == num_beams:
                cur_beam_indices = beam_indices
                cur_code_indices = code_indices
            else:
                cur_beam_indices = torch.div(torch.arange(
                    prev_nbeam * self.subvector_cents, device=doc_emb.device), self.subvector_cents, rounding_mode='floor')
                cur_code_indices = torch.arange(
                    self.subvector_cents, device=doc_emb.device).repeat(prev_nbeam)
            ncandidate = proba.size(1)
            assert len(cur_beam_indices) == len(
                cur_code_indices) == ncandidate
            # (bs, beam)
            if num_beams < ncandidate:
                if do_sample:
                    topk_indices = torch.multinomial(
                        proba, num_samples=num_beams)
                else:
                    _, topk_indices = proba.topk(num_beams, dim=-1)
                topk_prev_beams = cur_beam_indices[topk_indices].unsqueeze(-1)
                topk_cur_code = cur_code_indices[topk_indices]
                beam_scores = proba.gather(1, topk_indices)
                # (bs, beam, nsub)
                temp_index = torch.cat([temp_index.gather(
                    1, topk_prev_beams.expand(-1, -1, temp_index.size(-1))), topk_cur_code.unsqueeze(-1)], dim=-1)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    # (bs, beam, dim)
                    temp_embed = temp_embed.gather(
                        1, topk_prev_beams.expand(-1, -1, temp_embed.size(-1))) - codebook[i][topk_cur_code][..., :self.last_dim]
            else:
                beam_scores = proba
                temp_index = torch.cat([temp_index.repeat_interleave(
                    self.subvector_cents, dim=1), cur_code_indices.unsqueeze(-1).unsqueeze(0).expand(batch_size, -1, -1)], dim=-1)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    temp_embed = temp_embed.repeat_interleave(
                        self.subvector_cents, dim=1) - codebook[i][cur_code_indices][..., :self.last_dim]
        assert beam_scores.size(1) == num_beams
        topk_label = temp_index[:, :, 1:]
        if return_proba:
            return topk_label, beam_scores
        else:
            return topk_label

    @torch.no_grad()
    def get_topk_document_mapping(
        self,
        doc_embeddings: np.ndarray,
        rank: int,
        nrank: int,
        num_return_sequences: int,
        batch_size: int = 1024,
    ):
        num_docs = doc_embeddings.shape[0]
        num_docs_per_worker = num_docs // nrank
        start = num_docs_per_worker * rank
        if rank + 1 == nrank:
            ending = num_docs
        else:
            ending = start + num_docs_per_worker
        part_num_docs = ending - start
        all_topk_label = torch.empty(
            (part_num_docs, num_return_sequences, self.subvector_num), dtype=torch.int32, device='cpu')
        for i in tqdm(range(start, ending, batch_size), desc=f'Document TopK Mapping {rank}'):
            batch_start = i
            batch_ending = min(i+batch_size, ending)
            cur_embeddings = torch.tensor(
                doc_embeddings[batch_start:batch_ending], device='cpu')
            all_topk_label[batch_start-start:batch_ending-start] = self.beam_search(
                cur_embeddings, num_return_sequences)
        return all_topk_label

    def get_reconstruct_loss_for_embeddings(self, embeddings, labels):
        batch_size = labels.shape[0]
        nidim = len(labels.shape)
        assert nidim == 2
        codebook = self.get_codebook()[..., :self.last_dim].unsqueeze(
            0).expand(labels.size(0), -1, -1, -1)
        expand_target = (-1,) * (nidim + 1) + (self.last_dim,)
        labels = labels.unsqueeze(-1).unsqueeze(-1).expand(*expand_target)
        vectors = codebook.gather(-2, index=labels).squeeze(-2)
        if self.pq_type == 'pq':
            diff = embeddings - vectors.view(batch_size, -1)
        elif self.pq_type == 'rq':
            diffs = []
            cur_embeddings = embeddings
            for i in range(self.subvector_num):
                cur_embeddings = cur_embeddings - vectors[:, i, :]
                diffs.append(cur_embeddings)
            diff = torch.stack(diffs, 1)
        else:
            assert False
        loss = (diff ** 2).mean()
        # current using mean on all
        # TODO: try sum in sample and mean cross samples
        return loss

    def get_reconstruct_vector(self, index, codebook=None):
        nidim = len(index.shape)
        assert nidim in (1, 2)
        if codebook is None:
            # assert nidim == 1
            codebook = self.get_codebook()[..., :self.last_dim]
        expand_target = (-1,) * (nidim + 1) + (self.last_dim,)
        vectors = codebook.gather(
            -2, index=index.unsqueeze(-1).unsqueeze(-1).expand(*expand_target)).squeeze(-2)
        if self.pq_type == 'rq':
            vectors = torch.sum(vectors, dim=-2)
        else:
            view_target = (index.shape[0], -1) if nidim == 2 else (-1,)
            vectors = vectors.view(view_target)
        if self.pq_type == 'opq':
            vectors = torch.matmul(vectors, self.rotate)
        return vectors

    def get_reconstruct_vector_matrix_multiply(self, index):
        # index: (bs, nsub, ncent); codebook: (nsub, ncent, last_dim)
        bs = index.shape[0]
        codebook = self.get_codebook()[..., :self.last_dim]
        index = index.view(-1, self.subvector_cents).unsqueeze(1)
        codebook = codebook.unsqueeze(0).expand(
            bs, -1, -1, -1).reshape(-1, self.subvector_cents, self.last_dim)
        output = torch.bmm(index, codebook).squeeze(
            1).view(bs, self.subvector_num, self.last_dim)
        if self.pq_type == 'rq':
            output = torch.sum(output, dim=1)
        else:
            output = output.view(bs, -1)
        return output

if __name__ == '__main__':
    # offline generate faiss index
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--embedding_path', type=str, required=True)
    parser.add_argument('--save_path', type=str, required=True)
    parser.add_argument('--dist_mode', type=str, default='l2')
    parser.add_argument('--pq_type', type=str, default='pq')
    parser.add_argument('--subvector_num', type=int, default=4)
    parser.add_argument('--subvector_bits', type=int, default=4)
    parser.add_argument('--dim', type=int, default=768)
    args = parser.parse_args()
    assert args.dist_mode in ('l2', 'ip', 'iptol2')
    assert args.pq_type in ('pq', 'opq', 'rq')

    doc_embeddings = np.memmap(
        args.embedding_path, dtype=np.float32, mode='r').reshape(-1, args.dim)

    pq = ProductQuantization(
        args.pq_type,
        args.subvector_num,
        args.subvector_bits,
        args.dist_mode,
        args.emb_size,
    )
    pq.build_faiss_index(doc_embeddings, args.save_path)


#User function Template for python3

class Solution:

    #Function to find the maximum number of meetings that can
    #be performed in a meeting room.
    def maximumMeetings(self,n,start,end):

        # T.C-> O(N) for adding + O(NlogN) for sorting + O(N) for Traversing,checking & adding
        # Ultimate it is -> O(NlogN)
        # S.C-> O(N)

         #creating new Array
        newArray=[]
        for i in range(n):
            newArray.append([start[i], end[i], i+1])

        # sorting newArray
        newArray.sort(key=lambda x : x[1])
        # print(newArray)

        # performing the actual operation
        newAnsArray=[]
        newAnsArray.append(newArray[0][2])
        limit=newArray[0][1]

        for i in range(1,len(newArray)):
            if newArray[i][0]>limit:
                limit = newArray[i][1]
                newAnsArray.append(newArray[i][2])
        return len(newAnsArray)

Geek for Geeks/GreedyApproachStriver/N meetings in one room or  Activity Selection.py


from pwn import u32, u64, p32, p64

def unpack(s):
    if 0 < len(s) <= 4:
        return u32(s.ljust(4, '\x00'))
    elif len(s) <= 8:
        return u64(s.ljust(8, '\x00'))

def pack(n):
    if 0 < n <= 0xffffffff:
        return p32(n)
    elif n <= 0xffffffffffffffff:
        return p64(n)


from argparse import Namespace
from logging import info, debug

from external.config.metadata import CfgMetadata
from business.core.report import Report

class Metric:

    def __init__(self, metadata: CfgMetadata):
        self.jira = None
        self.metadata = metadata

    def set_adapters(self, adapters: dict) -> None:
        self.jira = adapters.get('jira')

    def get_cli(self) -> dict:
        return dict(
            name='metric',
            help='Calculate a team metrics',
            aliases='m',
            option=[
                dict(
                    name='begin',
                    help='Begin date of the calculation period (inclusive): <YYYY-MM-DD>',
                ),
                dict(
                    name='end',
                    help='End date of the calculation period (inclusive): <YYYY-MM-DD>',
                ),
            ],
            func=self._calculate,
        )

    def _calculate(self, args: Namespace) -> None:
        assert self.jira
        assert self.metadata

        period = dict(
            begin = args.begin,
            end = args.end,
        )
        report = Report(self.jira, self.metadata.get_general())

        info('Calculate team metrics for a period "{begin}" - "{end}"'.format(**period))

        for team in self.metadata.get_teams():
            print('{}:'.format(team.get('name')))

            metrics = report.get_metrics(team, period)
            for metric, value in sorted(metrics.items()):
                #* fractional value in hours required
                print('- {metric}: {value:0.3f}'.format(
                    metric=metric,
                    value=self.__convert_second_to_hour(value),
                ))

    def __convert_second_to_hour(self, seconds: int) -> float:
        return seconds / 3600.0


from django.shortcuts import render, redirect
from django.http import HttpResponse
from .models import ApplicationForm
from django.views.generic import DetailView
from django.http import JsonResponse
from PIL import Image
from django.core import serializers
from django.contrib import messages
from django.core.paginator import Paginator,PageNotAnInteger,EmptyPage
from portal.views import login_required
import json
from datetime import datetime
from django.db.models import Q

# Create your views here.

@login_required
def formSave(request):
    if request.method == "POST":
        reg = request.POST['regno'] #regno name is same on html field
        try:
            im = request.FILES["image"]
        except:
            im = None
        is_ajax = request.headers.get('X-Requested-With') == 'XMLHttpRequest'
       # check if uploaded image is valid (for example not video file ) .
        if not im == None:
            try:
                Image.open(im)
            except:
                if is_ajax:
                    return JsonResponse({'status':'500','message':'Invalid image    !'})
                else:
                    # messages.error(request, 'sorry, your image is invalid')
                    return redirect('save')
        dt = request.POST['date']
        sname = request.POST['surname']
        nm = request.POST['name']
        fn = request.POST['fname']
        mn = request.POST['mname']
        ad = request.POST['adno']
        db=request.POST['dob']
        cty = request.POST['city']
        gen= request.POST['gender']
        pl= request.POST['place']
        dist= request.POST['dist']
        st= request.POST['State']
        phy= request.POST['physical problems']
        school= request.POST['schoolname']
        data = ApplicationForm(reg_no = reg,surname=sname,name=nm,fname=fn,mname=mn, adhar_no =ad,
        gender=gen,place = pl,dist=dist,state=st,physical=phy,schoolname=school,date=dt,dob=db,profile_pic=im,city=cty)
        data.save()
        return JsonResponse({'status':'200','message':'Saved successfully!'})
    return render(request,'admission-form.html')

@login_required
def Show_View(request):
    data = ApplicationForm.objects.all()
    return render(request, 'admission-form.html')

def Home_View(request):
    return render(request, 'login.html')

@login_required
def Datatable_View(request):
    data = ApplicationForm.objects.all()
    per_page= 5
    obj_data = Paginator(data,per_page) #it tells how much data show per page
    first_page = obj_data.page(1).object_list
    last_page = obj_data.page_range
    context = {
        'obj_data':obj_data,
        'first_page':first_page,
        'last_page':last_page,
        'data':data
    }
    if request.method == 'POST':
        #getting page number
        page_no = request.POST.get('page_no', None)
        results = list(obj_data.page(page_no).object_list.values('id','reg_no','date','surname','fname','name','adhar_no','mname','gender','place','city','dist','state','state','physical','schoolname',))
        return JsonResponse({"results":results})
    return render(request, 'data-tables.html',context)

@login_required
def search_field(request):
    if request.method =="POST":
        is_ajax = request.headers.get('X-Requested-With') == 'XMLHttpRequest'
        if is_ajax:
            nam= request.POST.get('search_text')
            name_list = ApplicationForm.objects.filter(name__startswith=nam).values()
            return JsonResponse({'status':'200','message':'Ok!','users':list(name_list)})
    return render(request,'search.html')

@login_required
def Student_details(request, id):
    data = ApplicationForm.objects.get(id=id)
    return render(request,'details.html', {'data':data})

@login_required
def Update_View(request):
    if request.method == "POST":
        try:
            im = request.FILES["image"]
        except:
            im = None

        id = request.POST.get('user_id')
        print(id)
        reg = request.POST['regno']
        dt = request.POST['date_reg']
        db=request.POST['dob']

        sname = request.POST['surname']
        nm = request.POST['name']
        fn = request.POST['fname']
        mn = request.POST['mname']
        ad = request.POST['adno']

        cty = request.POST['city']
        gen= request.POST['gender']
        print(gen)
        pl= request.POST['place']
        dist= request.POST['dist']
        st= request.POST['State']
        phy= request.POST['physical problems']
        school= request.POST['schoolname']

        # d = {"num_pages":40, author:"Jack", date:"3324"}
        updateRecords = {
            'reg_no' :reg,
            'surname':sname,
            'name':nm,
            'fname':fn,
            'mname':mn,
            'adhar_no' :ad,
            'gender':gen,
            'place' : pl,
            'dist':dist,
            'state':st,
            'physical':phy,
            'schoolname':school,
            'date':dt,
            'dob':db,
            'city':cty
        }

        if im !=None:
            updateRecords['profile_pic']=im

        user = ApplicationForm.objects.get(id = id)
        print(user)
        user.update(**updateRecords)
        return redirect('search')

def Edit_user_modal(request,id):
    data = ApplicationForm.objects.get(id=id)

    return render(request,'update.html',{'data':data})

def Delete_View(request):
    if request.method == "POST":
        id1 = request.POST.get(id)
        print(id1)
        data = ApplicationForm.objects.get(id = id1)
        data.delete()
        return HttpResponse('deleted')

def delete(request):

    if request.method == "POST":
        id1 = request.POST.get('id')

        data = ApplicationForm.objects.get(id = id1)
        print(data)
        data.delete()
        return redirect('search.html')
    return HttpResponse('deleted')

#JSON RESPONSE

def Data_JSON(request):

    search = str(request.GET["search[value]"])

    offset = int(request.GET['start'])

    limit = int(request.GET['length'])
    count   =   ApplicationForm.objects.count()

    records =   ApplicationForm.objects.all().order_by('name')[offset:limit+offset]

    limit = limit+offset

    # CODE FOR SORTING
    col = int(request.GET['order[0][column]'])
    sort = str(request.GET['order[0][dir]'])
    columns = ['id','name','surname']
    sortable_col = str(columns[col])

    if search =="":
        if sort=="asc":
            sorted_data = ApplicationForm.objects.order_by(sortable_col)[offset:limit+offset]

        else:
            sorted_data = ApplicationForm.objects.order_by(sortable_col).reverse()[offset:limit+offset]

        # all_data =  [records.get_data() for records in records ]
        all_data =  [sorted_data.get_data() for sorted_data in sorted_data ]
        print("You are in if")

        all_records = {
            "recordsTotal": count,
            "recordsFiltered":count,
            "data": all_data,
         }

        return JsonResponse(all_records,safe=False,content_type = "application/json")
    else:

        data = ApplicationForm.objects.filter(Q(name__icontains=search) | Q(surname__icontains=search)).order_by('name')[offset:limit]
        q_data = [data.get_data() for data in data ]

        records = {
            "recordsTotal": count,    # The total count of records in your database, if query exist, it is the count of filtered records
            "recordsFiltered":count,
            "data": q_data,    # Main data to be rendered in table
         }

        return JsonResponse(records,safe=False,content_type = "application/json")

    # else:
    #     print("blank")
    #             # data = records.get_queryset().values()
    #     all_data =  [records.get_data() for records in records ]
    #     print(all_data)

    #     all_records = {
    #         "recordsTotal": count,
    #         "recordsFiltered":count,
    #         "data": all_data,
    #      }

    #     return JsonResponse(all_records,safe=False,content_type = "application/json")

#HTML RENDER

def Data_table_view(request):
    return render(request,'data_tab.html')

def sorting(col,sort):

    pass

def Search_in_Datatable(request):

    return HttpResponse("<h5> Hello </h5> ")


import json

from flask import Flask, jsonify, render_template, request

import predict

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict-next-words', methods=["POST"])
def predict_next_words():
    body = request.json['body']
    language = request.json['language']
    res = predict.predict_next_words(body, language)

    return jsonify({'nextWords': res})

@app.route('/predict-language', methods=["POST"])
def predict_language():
    body = request.json['body']
    res = predict.predict_language(body)

    return jsonify({'language': res})

if __name__ == '__main__':
    app.run()


#!/usr/bin/env python
import re, os, sys, shutil
from math import *
from string import *
from optparse import OptionParser
import operator

def normalize_tag_count(input_file, ColumnIndex, total, output_file):
	"""
	Column index is 0 based
	"""
	infile = open(input_file,'r')
	outfile = open(output_file, 'w')
	for line in infile:
		if not re.match("#", line):
			line = line.strip()
			sline = line.split()
			if len(sline) >= (ColumnIndex + 1):
				sline[ColumnIndex] = str(atof(sline[ColumnIndex])/total)
				outline = '\t'.join(sline) +'\n'
				outfile.write(outline)
	infile.close()
	outfile.close()

def total_counts(file, column):
	"""
	Column index is 0 based
	"""
	total = 0.0
	infile = open(file,'r')
	for line in infile:
		if not re.match("#", line):
			line = line.strip()
			sline = line.split()
			if len(sline) >= (column + 1):
				total += atof(sline[column])
	infile.close()
	return total

def main(argv):
	parser = OptionParser()
	parser.add_option("-i", "--inputfile", action="store", type="string", dest="input_file", metavar="<file>", help="original data file to be normalized")
	parser.add_option("-a", "--tagcountcolumn", action="store", type="int", dest="colum", metavar="<int>", help="colum index for tag counts to be normalized, start from 0")
	parser.add_option("-t", "--ScalingFactor", action="store", type="float", dest="ScalingFactor", metavar="<float>", help="Scaling factor for normalization, for example, scaling factor is 1000000 if normalization by total read count per million")
	parser.add_option("-o", "--outputfile", action="store", type="string", dest="output_file", metavar="<file>", help="output file name")

	(opt, args) = parser.parse_args(argv)
	if len(argv) < 8:
        	parser.print_help()
        	sys.exit(1)

	total = total_counts(opt.input_file, opt.colum) / opt.ScalingFactor
	normalize_tag_count(opt.input_file, opt.colum, total, opt.output_file)

if __name__ == "__main__":
	main(sys.argv)


from __future__ import print_function
import ConfigParser, os
from sys import argv

from apiclient import discovery
import auth
import settings
from auth import Authorize

from apiclient import discovery
from apiclient.http import MediaFileUpload

#Command arguments to test functionality, should be passed via torrent in future
script, filename = argv

http = Authorize()

serviceDrive = discovery.build('drive', 'v2', http=http)

#Function for uploading filename set in argv when executed
def main(filename):
    FilePath = os.path.abspath(filename)
    body = {
            'title': filename
    }
    body['parents'] = [{'id' : settings.googledrivedir}]
    media = MediaFileUpload(FilePath, chunksize=settings.chunksize, resumable=True)
    response = serviceDrive.files().insert(body=body, media_body=media).execute()
    if settings.nonDefaultPermissions == True:
        fileID = response['id']
        setPermissions(fileID)
    print(response['alternateLink'])

def setPermissions(file_id):
    service = serviceDrive
    newPermissions = {
        'value': settings.permissionValue,
        'type': settings.permissionType,
        'role': settings.permissionRole
    }
    return service.permissions().insert(fileId=file_id,
        body=newPermissions).execute()

if __name__ == '__main__':
    main(filename)