import json
from pyspark.sql import SparkSession
from pyspark.sql.functions import concat, col, lit, split, to_date, date_format
import os
import time

# Create a SparkSession
spark = SparkSession.builder.getOrCreate()

# Start the timer
start_time = time.time()

# Load the config file
with open('config.json') as f:
    config = json.load(f)

#Check if the input_csv file exists and is in csv format
if 'input_csv' not in config or not config['input_csv'].endswith('.csv') or not os.path.exists(config['input_csv']):
    print("No CSV file is selected or the file is invalid.")
    exit()

# Load the DataFrame from the CSV file
df = spark.read.csv(config['input_csv'], header=True, inferSchema=True)

# Convert the DataFrame to a Parquet file
df.write.mode("overwrite").parquet('temp.parquet')
# Load the DataFrame from the Parquet file
df = spark.read.parquet('temp.parquet')

# Apply transformations
for transformation in config['transformations']:
    if transformation['type'] == 'date_format':
        df = df.withColumn(transformation['column'], to_date(col(transformation['column']), transformation['input_format']))
        df = df.withColumn(transformation['column'], date_format(col(transformation['column']), transformation['output_format']))
    elif transformation['type'] == 'concat':
        df = df.withColumn(transformation['output_column'], concat(*[col(c) for c in transformation['columns']], lit(transformation['separator'])))
    elif transformation['type'] == 'split':
        split_col = split(df[transformation['column']], transformation['separator'])
        for i, output_column in enumerate(transformation['output_columns']):
            df = df.withColumn(output_column, split_col.getItem(i))
    elif transformation['type'] == 'drop':
        df = df.drop(transformation['column'])

# Get the total number of records
total_records = df.count()

# Print the total number of records
print("Total records: ", total_records)

# Save the transformed DataFrame to a new Parquet file
df.write.mode('overwrite').parquet(config['output_parquet'])
df = spark.read.parquet(config['output_parquet'])
df.show()

# End the timer
end_time = time.time()

# Calculate the elapsed time
elapsed_time = end_time - start_time

# Print the elapsed time
print("Time taken: {} seconds".format(elapsed_time))


'''
Created on 06.12.2013

@author: hfrieden

Import an Arma 2/Arma 3 unbinarized MDL file

'''
import struct
import bpy
import bmesh
import os.path as path
import ArmaToolbox
import ArmaTools

def getLayerMask(layer):
    res = [False, False, False, False, False,
           False, False, False, False, False,
           False, False, False, False, False,
           False, False, False, False, False]
    res[layer % 20] = True
    return res

# Datatype reading
def readULong(filePtr):
    return struct.unpack("i", filePtr.read(4))[0]

def readSignature(filePtr):
    return filePtr.read(4)

def readFloat(filePtr):
    return struct.unpack("f", filePtr.read(4))[0]

def readChar(filePtr):
    return struct.unpack("c", filePtr.read(1))[0]

def readByte(filePtr):
    return struct.unpack("b", filePtr.read(1))[0]

def readString(filePtr):
    res = b''
    t = True
    while t:
        a = filePtr.read(1)
        if a != b'\000':
            res = res + a
        else:
            t = False

    return res.decode("utf-8")

def makeLodName(fileName, lodLevel):
    lodName = path.basename(fileName)
    lodName = lodName.split(".")[0]
    lodName = "{0}.{1}".format(lodName, lodLevel)
    return lodName

def maybeAddEdgeSplit(obj):
    obj.data.use_auto_smooth = True
    obj.data.auto_smooth_angle = 3.1415927

    #modifier = obj.modifiers.get("FHQ_ARMA_Toolbox_EdgeSplit")
    #if modifier is None:
    #    modifier = obj.modifiers.new("FHQ_ARMA_Toolbox_EdgeSplit",
    #             type='EDGE_SPLIT')
    #
    #    modifier.show_expanded = False
    #    modifier.use_edge_angle = False # Want only sharp edges
    #    modifier.use_edge_sharp = True

    #obj.data.show_edge_sharp = True

def correctedResolution(r):
    res = int(r)

    if (r < 1000):
        return r #res

    values ={
        10000000000000 : 'G',
        3000000000000000 : 'RD',
        #11010 : 'SV2',
        8000000000000000 : 'VCG',
        10000 : 'S1',
        14000000000000000 : 'VPFG',
        17000000000000000 : 'SP',
        #10010 : 'S2',
        12000000000000000 : 'VCFG',
        20000000000000000 : 'SVVG',
        1200 : 'VC',
        9000000000000000 : 'VCFG',
        15000000000000000 : 'VGG',
        13000000000000000 : 'VPG',
        18000000000000000 : 'SVVC',
        1000000000000000 : 'M',
        1100 : 'P',
        21000000000000000 : 'WRL',
        4000000000000000 : 'PTH',
        40000000000000 : 'GPX',
        7000000000000000 : 'FG',
        10000000000000000 : 'VC',
        6000000000000000 : 'VG',
        1000 : 'GN',
        16000000000000000 : 'VGFG',
        20000000000000 : 'GB',
        19000000000000000 : 'SVVP',
        2000000000000000 : 'LC',
        11000 : 'SV1',
        20000 : 'ED',
        5000000000000000 : 'HP',
        11000000000000000 : 'VCG',
    }
    error = 1000000000000000000000
    value = -1
    for n in values:
        x = abs(n-res)
        if x < error:
            error = x
            value = n
    return value

def resolutionName(r):
    res = int(r)

    if (r < 1000):
        return str(res)

    values ={
        1.000e+3:'View Gunner',
        1.100e+3:'View Pilot',
        1.200e+3:'View Cargo',
        1.000e+4:'Stencil Shadow',
        2.000e+4:'Edit',
        #1.001e+4:'Stencil Shadow 2',
        1.100e+4:'Shadow Volume',
        #1.101e+4:'Shadow Volume 2',
        1.000e+13:'Geometry',
        1.000e+15:'Memory',
        2.000e+15:'Land Contact',
        3.000e+15:'Roadway',
        4.000e+15:'Paths',
        5.000e+15:'Hit Points',
        6.000e+15:'View Geometry',
        7.000e+15:'Fire Geometry',
        8.000e+15:'View Cargo Geometry',
        9.000e+15:'View Cargo Fire Geometry',
        1.000e+16:'View Commander',
        1.100e+16:'View Commander Geometry',
        1.200e+16:'View Commander Fire Geometry',
        1.300e+16:'View Pilot Geometry',
        1.400e+16:'View Pilot Fire Geometry',
        1.500e+16:'View Gunner Geometry',
        1.600e+16:'View Gunner Fire Geometry',
        1.700e+16:'Sub Parts',
        1.800e+16:'Cargo View shadow volume',
        1.900e+16:'Pilot View shadow volume',
        2.000e+16:'Gunner View shadow volume',
        2.100e+16:'Wreckage',
        2.000e+13:'Geometry Buoyancy',
        4.000e+13:'Geometry PhysX'
    }
    error = 1000000000000000000000
    value = -1
    for n in values:
        x = abs(n-res)
        if x < error:
            error = x
            value = n
    ret = values.get(value, "?")
    if value == 1.000e+4 or value == 2.000e+4:
        ret = ret + " " + str(r-value)

    return ret

def decodeWeight(b):
    if  b == 0:
        return  0.0
    elif b == 2:
        return 1.0
    elif b > 2:
        return 1.0 - round( (b-2) / 2.55555 )*0.01
    elif b < 0:
        return -round( b / 2.55555 ) * 0.01
    else:
        # print ("decodeWeight(",b,") returns 1.0 as else case")
        return 1.0 #TODO: Correct?

def loadLOD(context, filePtr, objectName, materialData, layerFlag, lodnr):
    global objectLayers
    meshName = objectName
    weightArray = []
    # Check for P3DM signature
    sig = readSignature(filePtr)
    if sig != b'P3DM':
        return -1

    # Read major and minor version
    major = readULong(filePtr)
    minor = readULong(filePtr)
    if major != 0x1c:
        print("Unknown major version {0}".format(major))
        return -1
    if minor != 0x100:
        print("Unknown minor version {0}".format(minor))
        return -1

    numPoints   = readULong(filePtr)
    numNormals  = readULong(filePtr)
    numFaces    = readULong(filePtr)
    print("read lod")
    dummyFlags  = readULong(filePtr)

    # Read the Points. Points are XYZTriples followed by an ULONG flags word
    verts = []
    for i in range(0, numPoints):
        point = struct.unpack("fffi", filePtr.read(16))
        pnt = [point[0],  point[2], point[1]]
        verts.append(pnt)

    print("normals (",numNormals, ")...")
    normals = []
    for i in range(0, numNormals):
        normal = struct.unpack("fff", filePtr.read(12))
        nrm = [normal[0], normal[1], normal[2]]
        normals.append(normal)
        #print ("Normal = ", normal)

    faceData = []
    faces = []
    print("faces...")
    # Start reading and adding faces
    for i in range(0, numFaces):
        numSides = readULong(filePtr)
        # Vertex table
        vIdx = []
        nrmIdx = []
        uvs = []
        for n in range(0, 4):
            vtable = struct.unpack("iiff", filePtr.read(16))
            if n<numSides:
                vIdx.append(vtable[0])
                nrmIdx.append(vtable[1])
                uvs.append( [vtable[2], vtable[3]])
        faceFlags = readULong(filePtr)
        textureName = readString(filePtr)
        materialName = readString(filePtr)
        faceData.append(
            (numSides, nrmIdx, uvs, faceFlags, textureName, materialName)
        )
        faces.append(vIdx)
        # Handle the material if it doesn't exists yet
        if len(textureName) > 0 or  len(materialName)>0:
            try:
                materialData[(textureName, materialName)]
            except:
                # Need to create a new material for this
                #mat =  bpy.data.materials.new("Arma Material")
                mat = bpy.data.materials.new(path.basename(textureName) + " :: " + path.basename(materialName))
                mat.armaMatProps.colorString = textureName
                mat.armaMatProps.rvMat   = materialName
                if len(textureName) > 0 and textureName[0] == '#':
                    mat.armaMatProps.texType = 'Custom'
                    mat.armaMatProps.colorString = textureName
                else:
                    mat.armaMatProps.texType = 'Texture'
                    mat.armaMatProps.texture = textureName
                    mat.armaMatProps.colorString = ""

                materialData[(textureName, materialName)] = mat

    if readSignature(filePtr) != b'TAGG':
        print("No tagg signature")
        return -1;

    # Create the mesh. Doing it here makes the named selections
    # easier to read.
    mymesh = bpy.data.meshes.new(name=meshName)
    mymesh.from_pydata(verts, [], faces)

    mymesh.update(calc_edges = True)

    obj = bpy.data.objects.new(meshName, mymesh)

    # TODO: Maybe add a "logical Collection" option that
    # Collects all geometries, shadows, custom etc in a collection.

    scn = bpy.context.scene

    coll = bpy.data.collections.new(meshName)
    context.scene.collection.children.link(coll)
    coll.objects.link(obj)

    #NEIN! coll.hide_viewport = True
    #scn.objects.link(obj)
    #scn.objects.active = obj

    # Build Edge database to make finding sharp edges easier
    edgeDict = dict()
    for edge in mymesh.edges:
        v1 = edge.vertices[0]
        v2 = edge.vertices[1]
        if (v1 > v2): # Swap if out of order
            temp = v2
            v2 = v1
            v1 = temp
        #print(f"adding edge index {edge.index} as ({v1},{v2}) to edge dictionary")
        edgeDict[(v1,v2)] = edge.index

    print("taggs")
    loop = True
    sharpEdges = None
    weight = None
    while loop:
        active = readChar(filePtr)
        tagName = readString(filePtr)

        numBytes = readULong(filePtr)

        #print ("tagg: ",tagName, " size ", numBytes)

        if active == b'\000':
            if numBytes != 0:
                filePtr.seek(numBytes, 1)
        else:
            if tagName == "#EndOfFile#":
                loop = False
            elif tagName == "#SharpEdges#":
                # Read Sharp Edges
                sharpEdges = []
                for i in range(0,numBytes,8):
                    n1 = readULong(filePtr)
                    n2 = readULong(filePtr)
                    sharpEdges.append([n1, n2])
                #print ("sharp edges", sharpEdges)
            elif tagName == "#Property#":
                # Read named property
                propName  = struct.unpack("64s", filePtr.read(64))[0].decode("utf-8")
                propValue = struct.unpack("64s", filePtr.read(64))[0].decode("utf-8")
                item = obj.armaObjProps.namedProps.add()
                item.name=propName;
                item.value=propValue
            elif tagName == "#UVSet#":
                id = readULong(filePtr)
                layerName = "UVSet " + str(id)
                if id == 0:
                    # Name first layer "UVMap" so that there isn't any fuckups with uv sets
                    layerName = "UVMap"
                #print("adding UV set " + layerName)
                mymesh.uv_layers.new(name=layerName)
                layer = mymesh.uv_layers[-1]
                index = 0
                for faceIdx in range(0,numFaces):
                    n = faceData[faceIdx][0]
                    for x in range(0,n):
                        u = readFloat(filePtr)
                        v = readFloat(filePtr)
                        layer.data[index].uv = [u,1 - v]
                        index += 1
            elif tagName == "#Mass#":
                weightArray = []
                weight = 0;
                for idx in range (0,numPoints):
                    f = readFloat(filePtr)
                    weightArray.append(f)
                    weight += f
            elif tagName[0] == '#':
                # System tag we don't read
                filePtr.seek(numBytes, 1)
            else:
                # Named Selection
                # Add a vertex group
                # First, check the tagName for a proxy
                newVGrp = True
                if len(tagName) > 5:
                    if tagName[:6] == "proxy:":
                        newVGrp = False
                        vgrp = obj.vertex_groups.new(name = "@@armaproxy")
                        prp = obj.armaObjProps.proxyArray
                        prx = tagName.split(":")[1]
                        if prx.find(".") != -1:
                            a = prx.split(".")
                            prx = a[0]
                            idx = a[-1]
                            if len(idx) == 0:
                                idx = "1"
                        else:
                            idx = "1"
                        n = prp.add()
                        n.name = vgrp.name
                        n.index = int(idx)
                        n.path = "P:" + prx
                        tagName = "@@armyproxy"

                if newVGrp == True:
                    vgrp = obj.vertex_groups.new(name = tagName)
                for i in range(0, numPoints):
                    b = readByte(filePtr)
                    w = decodeWeight(b)
                    if (w>0):
                        vgrp.add([i],float(w),'REPLACE')
                    #print("b = ",b,"w = ", w)
                for i in range(0, numFaces):
                    b = readByte(filePtr)
                    w = decodeWeight(b)
                #    print("b = ",b,"w = ", w)
                #    if w== 1.0:
                #        pPoly = obj.data.polygons[i]
                #        for n in range(0,len(pPoly.vertices)):
                #            idx = pPoly.vertices[n]
                #            vgrp.add([idx], w, 'REPLACE')
                #filePtr.seek(numFaces, 1)

    # Done with the taggs, only the resolution is left to read
    resolution = readFloat(filePtr)
    #meshName = meshName + "." + resolutionName(resolution)
    meshName = resolutionName(resolution)
    mymesh.name = meshName
    obj.name = meshName
    coll.name = meshName

    print("materials...")
    indexData = {}
    # Set up materials
    for faceIdx in range(0,numFaces):
        fd = faceData[faceIdx]

        textureName = fd[4]
        materialName = fd[5]

        try:
            mat = materialData[(textureName, materialName)]
            # Add the material if it isn't in
            if mat.name not in mymesh.materials:
                mymesh.materials.append(mat)
                thisMatIndex = len(mymesh.materials)-1
                indexData[mat] = thisMatIndex
                #print("added new material at " + str(thisMatIndex))
            else:
                thisMatIndex = indexData [mat]
                #print("old material " + str(thisMatIndex))

            mymesh.polygons[faceIdx].material_index = thisMatIndex
        except:
            pass

    print("sharp edges")
    # Set sharp edges
    #if sharpEdges is not None:
    #    for edge in mymesh.edges:
    #        v1 = edge.vertices[0]
    #        v2 = edge.vertices[1]
    #        if [v1,v2] in sharpEdges:
    #            mymesh.edges[edge.index].use_edge_sharp = True
    #        elif [v2,v1] in sharpEdges:
    #            mymesh.edges[edge.index].use_edge_sharp = True
    #        else:
    #            print(f"Edge pair {v1},{v2} not found in edges")

    # New Code
    if sharpEdges is not None:
        for sharpEdge in sharpEdges:
            v1 = sharpEdge[0]
            v2 = sharpEdge[1]
            if (v1 > v2): # Swap if out of order
                temp = v2
                v2 = v1
                v1 = temp
            try: # Apparently, some models have sharp edges that (no longer) exist.
                idx = edgeDict[(v1,v2)]
                mymesh.edges[idx].use_edge_sharp = True
            except:
                print(f"WARNING: Edge {v1},{v2} does not exist")

    #for pair in sharpEdges:
    #    p1 = pair[0]
    #    p2 = pair[1]
    #    edge = mymesh.edges.get([mymesh.vertices[p1], mymesh.vertices[p2]])
    #    print("edge = ", edge)
    #    #if edge != None:
    #    #    edge.use_edge_sharp = True

    # TODO: This causes faces with the same vertices but different normals to
    # be discarded. Don't want that
    #mymesh.validate()
    print("Normal calculation")
    mymesh.calc_normals()
    for poly in mymesh.polygons:
        poly.use_smooth = True

    print("Add edge split")
    maybeAddEdgeSplit(obj)
    #scn.update()
    obj.select_set(True)

    #if layerFlag == True:
    #    # Move to layer
    #    objectLayers = getLayerMask(lodnr)
    #    bpy.ops.object.move_to_layer(layers=objectLayers)

    hasSet = False
    oldres = resolution
    resolution = correctedResolution(resolution)
    offset = oldres - resolution
    obj.armaObjProps.isArmaObject = True
    if resolution <= 1000:
        obj.armaObjProps.lodDistance = resolution
        hasSet = True
    else:
        obj.armaObjProps.lodDistance = offset #0.0

    print("set LOD type")
    # Set the right LOD type
    lodPresets = ArmaToolbox.lodPresets

    for n in lodPresets:
        if float(n[0]) == resolution:
            obj.armaObjProps.lod = n[0]
            hasSet = True

    if hasSet == False:
        print("Error: unknown lod %f" % (resolution))
        print("resolution %d" % (correctedResolution(resolution)))

    print("weight")
    if weight is not None:
        obj.armaObjProps.mass = weight

    if len(weightArray) > 0:
        bm = bmesh.new()
        bm.from_mesh(obj.data)
        bm.verts.ensure_lookup_table()

        weight_layer = bm.verts.layers.float.new('FHQWeights')
        weight_layer = bm.verts.layers.float['FHQWeights']
        print(weight_layer)
        for i in range(0,len(weightArray)):
            bm.verts[i][weight_layer] = weightArray[i]

        bm.to_mesh(obj.data)

    obj.select_set(False)

    if obj.armaObjProps.lod == '1.000e+13' or obj.armaObjProps.lod == '4.000e+13':
        ArmaTools.attemptFixMassLod (obj)

    if obj.armaObjProps.lod == '-1.0':
        ArmaTools.PostProcessLOD(obj)

    print("done reading lod")
    return 0

# Main Import Routine
def importMDL(context, fileName, layerFlag):
    global objectLayers
    objectLayers = [True, False, False, False, False,
               False, False, False, False, False,
               False, False, False, False, False,
               False, False, False, False, False]
    currentLayer = 0

    filePtr = open(fileName, "rb")

    objName = path.basename(fileName).split(".")[0]

    # This is used to collect combinations of texture and rvmat
    # in order to generate Materials
    materialData = {}

    # Read the header
    sig = readSignature(filePtr)
    version = readULong(filePtr)
    numLods = readULong(filePtr)

    print ("Signature = {0}, version={1}, numLods = {2}".format(sig, version, numLods))

    if version != 257 or sig != b'MLOD':
        return -1

    # Start loading lods
    for i in range(0,numLods):
        if loadLOD(context, filePtr, objName, materialData, layerFlag, i) != 0:
            return -2

    filePtr.close()

    return 0


import pytest
from sqlobject import DatabaseIndex, ForeignKey, IntCol, MultipleJoin, \
    SQLObject, StringCol
from sqlobject.dberrors import DatabaseError, IntegrityError, \
    OperationalError, ProgrammingError
from sqlobject.tests.dbtest import raises, setupClass, supports

########################################
# Indexes
########################################

class SOIndex1(SQLObject):
    name = StringCol(length=100)
    number = IntCol()

    nameIndex = DatabaseIndex('name', unique=True)
    nameIndex2 = DatabaseIndex(name, number)
    nameIndex3 = DatabaseIndex({'column': name,
                                'length': 3})

class SOIndex2(SQLObject):
    name = StringCol(length=100)
    nameIndex = DatabaseIndex({'expression': 'lower(name)'})

def test_indexes_1():
    setupClass(SOIndex1)
    n = 0
    for name in 'blah blech boring yep yort snort'.split():
        n += 1
        SOIndex1(name=name, number=n)
    mod = SOIndex1._connection.module
    raises(
        (mod.ProgrammingError, mod.IntegrityError,
         mod.OperationalError, mod.DatabaseError,
         ProgrammingError, IntegrityError, OperationalError, DatabaseError),
        SOIndex1, name='blah', number=0)

def test_indexes_2():
    if not supports('expressionIndex'):
        pytest.skip("expressionIndex isn't supported")
    setupClass(SOIndex2)
    SOIndex2(name='')

class PersonIndexGet(SQLObject):
    firstName = StringCol(length=100)
    lastName = StringCol(length=100)
    age = IntCol(alternateID=True)
    nameIndex = DatabaseIndex(firstName, lastName, unique=True)

def test_index_get_1():
    setupClass(PersonIndexGet, force=True)

    PersonIndexGet(firstName='Eric', lastName='Idle', age=62)
    PersonIndexGet(firstName='Terry', lastName='Gilliam', age=65)
    PersonIndexGet(firstName='John', lastName='Cleese', age=66)

    PersonIndexGet.get(1)
    PersonIndexGet.nameIndex.get('Terry', 'Gilliam')
    PersonIndexGet.nameIndex.get(firstName='John', lastName='Cleese')

    raises(Exception, PersonIndexGet.nameIndex.get,
           firstName='Graham', lastName='Chapman')

    raises(Exception, PersonIndexGet.nameIndex.get,
           'Terry', lastName='Gilliam')

    raises(Exception, PersonIndexGet.nameIndex.get, 'Terry', 'Gilliam', 65)

    raises(Exception, PersonIndexGet.nameIndex.get, 'Terry')

class PersonIndexGet2(SQLObject):
    name = StringCol(alternateID=True, length=100)
    age = IntCol()
    addresses = MultipleJoin('AddressIndexGet2')

class AddressIndexGet2(SQLObject):
    person = ForeignKey('PersonIndexGet2', notNone=True)
    type = StringCol(notNone=True, length=100)
    street = StringCol(notNone=True)
    pk = DatabaseIndex(person, type, unique=True)

def test_index_get_2():
    setupClass([PersonIndexGet2, AddressIndexGet2])

    p = PersonIndexGet2(name='Terry Guilliam', age=64)
    AddressIndexGet2(person=p, type='home', street='Terry Street 234')
    AddressIndexGet2(person=p, type='work', street='Guilliam Street 234')

    AddressIndexGet2.pk.get(p, 'work')
    AddressIndexGet2.pk.get(person=p, type='work')


import csv
import operator
import decimal
import numpy as np
import matplotlib.pyplot as pl
import calendar
import sys
import os.path

h = open('../results/MonthlyTurnoverOldCerrado.csv','rb')
data = csv.reader(h)

months = []
bints = []
seasons = []
boundarys = []
sumprecips = []
avghumid = []
avgvisits = []

for column in data:
  months.append(column[0])
  bints.append(column[1])
  seasons.append(column[2])
  boundarys.append(column[3])
  sumprecips.append(column[4])
  avghumid.append(column[5])
  avgvisits.append(column[5])

h.close()

#remove header
for i in [months, bints, seasons, boundarys, sumprecips, avghumid, avgvisits]:
  del i[0]

pl.plot(sumprecips, avgvisits, 'ro')

pl.grid(True)

pl.title('sumprecip-Avgvisits')

plotpath = '../results/' + 'sumprecip-Avgvisits' + '.pdf'
pl.savefig(plotpath)

pl.show()


from contextlib import contextmanager
import sys
import os
import tempfile
from shutil import rmtree
from os import getcwd, chdir
from os.path import join, basename, dirname, isdir, abspath, sep
import unittest

import six
from six.moves import reload_module

from pylint import config, lint
from pylint.lint import PyLinter, Run, preprocess_options, \
     ArgumentPreprocessingError
from pylint.utils import MSG_STATE_SCOPE_CONFIG, MSG_STATE_SCOPE_MODULE, MSG_STATE_CONFIDENCE, \
    MessagesStore, PyLintASTWalker, MessageDefinition, FileState, \
    build_message_def, tokenize_module, UnknownMessage
from pylint.testutils import TestReporter, catch_warnings
from pylint.reporters import text, html
from pylint import checkers
from pylint.checkers.utils import check_messages
from pylint import interfaces

if os.name == 'java':
    if os._name == 'nt':
        HOME = 'USERPROFILE'
    else:
        HOME = 'HOME'
else:
    if sys.platform == 'win32':
        HOME = 'USERPROFILE'
    else:
        HOME = 'HOME'

@contextmanager
def fake_home():
    folder = tempfile.mkdtemp('fake-home')
    old_home = os.environ.get(HOME)
    try:
        os.environ[HOME] = folder
        yield
    finally:
        os.environ.pop('PYLINTRC', '')
        if old_home is None:
            del os.environ[HOME]
        else:
            os.environ[HOME] = old_home
        rmtree(folder, ignore_errors=True)

def remove(file):
    try:
        os.remove(file)
    except OSError:
        pass

HERE = abspath(dirname(__file__))
INPUTDIR = join(HERE, 'input')

@contextmanager
def tempdir():
    """Create a temp directory and change the current location to it.

    This is supposed to be used with a *with* statement.
    """
    tmp = tempfile.mkdtemp()

    # Get real path of tempfile, otherwise test fail on mac os x
    current_dir = getcwd()
    chdir(tmp)
    abs_tmp = abspath('.')

    try:
        yield abs_tmp
    finally:
        chdir(current_dir)
        rmtree(abs_tmp)

def create_files(paths, chroot='.'):
    """Creates directories and files found in <path>.

    :param paths: list of relative paths to files or directories
    :param chroot: the root directory in which paths will be created

    >>> from os.path import isdir, isfile
    >>> isdir('/tmp/a')
    False
    >>> create_files(['a/b/foo.py', 'a/b/c/', 'a/b/c/d/e.py'], '/tmp')
    >>> isdir('/tmp/a')
    True
    >>> isdir('/tmp/a/b/c')
    True
    >>> isfile('/tmp/a/b/c/d/e.py')
    True
    >>> isfile('/tmp/a/b/foo.py')
    True
    """
    dirs, files = set(), set()
    for path in paths:
        path = join(chroot, path)
        filename = basename(path)
        # path is a directory path
        if filename == '':
            dirs.add(path)
        # path is a filename path
        else:
            dirs.add(dirname(path))
            files.add(path)
    for dirpath in dirs:
        if not isdir(dirpath):
            os.makedirs(dirpath)
    for filepath in files:
        open(filepath, 'w').close()

class SysPathFixupTC(unittest.TestCase):
    def setUp(self):
        self.orig = list(sys.path)
        self.fake = [1, 2, 3]
        sys.path[:] = self.fake

    def tearDown(self):
        sys.path[:] = self.orig

    def test_no_args(self):
        with lint.fix_import_path([]):
            self.assertEqual(sys.path, self.fake)
        self.assertEqual(sys.path, self.fake)

    def test_one_arg(self):
        with tempdir() as chroot:
            create_files(['a/b/__init__.py'])
            expected = [join(chroot, 'a')] + self.fake

            cases = (
                ['a/b/'],
                ['a/b'],
                ['a/b/__init__.py'],
                ['a/'],
                ['a'],
            )

            self.assertEqual(sys.path, self.fake)
            for case in cases:
                with lint.fix_import_path(case):
                    self.assertEqual(sys.path, expected)
                self.assertEqual(sys.path, self.fake)

    def test_two_similar_args(self):
        with tempdir() as chroot:
            create_files(['a/b/__init__.py', 'a/c/__init__.py'])
            expected = [join(chroot, 'a')] + self.fake

            cases = (
                ['a/b', 'a/c'],
                ['a/c/', 'a/b/'],
                ['a/b/__init__.py', 'a/c/__init__.py'],
                ['a', 'a/c/__init__.py'],
            )

            self.assertEqual(sys.path, self.fake)
            for case in cases:
                with lint.fix_import_path(case):
                    self.assertEqual(sys.path, expected)
                self.assertEqual(sys.path, self.fake)

    def test_more_args(self):
        with tempdir() as chroot:
            create_files(['a/b/c/__init__.py', 'a/d/__init__.py', 'a/e/f.py'])
            expected = [
                join(chroot, suffix)
                for suffix in [sep.join(('a', 'b')), 'a', sep.join(('a', 'e'))]
            ] + self.fake

            cases = (
                ['a/b/c/__init__.py', 'a/d/__init__.py', 'a/e/f.py'],
                ['a/b/c', 'a', 'a/e'],
                ['a/b/c', 'a', 'a/b/c', 'a/e', 'a'],
            )

            self.assertEqual(sys.path, self.fake)
            for case in cases:
                with lint.fix_import_path(case):
                    self.assertEqual(sys.path, expected)
                self.assertEqual(sys.path, self.fake)

class PyLinterTC(unittest.TestCase):

    def setUp(self):
        self.linter = PyLinter()
        self.linter.disable('I')
        self.linter.config.persistent = 0
        # register checkers
        checkers.initialize(self.linter)
        self.linter.set_reporter(TestReporter())

    def init_linter(self):
        linter = self.linter
        linter.open()
        linter.set_current_module('toto')
        linter.file_state = FileState('toto')
        return linter

    def test_pylint_visit_method_taken_in_account(self):
        class CustomChecker(checkers.BaseChecker):
            __implements__ = interfaces.IAstroidChecker
            name = 'custom'
            msgs = {'W9999': ('', 'custom', '')}

            @check_messages('custom')
            def visit_class(self, _):
               pass

        self.linter.register_checker(CustomChecker(self.linter))
        self.linter.open()
        out = six.moves.StringIO()
        self.linter.set_reporter(text.TextReporter(out))
        self.linter.check('abc')

    def test_enable_message(self):
        linter = self.init_linter()
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('W0102'))
        linter.disable('W0101', scope='package')
        linter.disable('W0102', scope='module', line=1)
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertFalse(linter.is_message_enabled('W0102', 1))
        linter.set_current_module('tutu')
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('W0102'))
        linter.enable('W0101', scope='package')
        linter.enable('W0102', scope='module', line=1)
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('W0102', 1))

    def test_enable_message_category(self):
        linter = self.init_linter()
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('C0202'))
        linter.disable('W', scope='package')
        linter.disable('C', scope='module', line=1)
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('C0202'))
        self.assertFalse(linter.is_message_enabled('C0202', line=1))
        linter.set_current_module('tutu')
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('C0202'))
        linter.enable('W', scope='package')
        linter.enable('C', scope='module', line=1)
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('C0202'))
        self.assertTrue(linter.is_message_enabled('C0202', line=1))

    def test_message_state_scope(self):
        class FakeConfig(object):
            confidence = ['HIGH']

        linter = self.init_linter()
        linter.disable('C0202')
        self.assertEqual(MSG_STATE_SCOPE_CONFIG,
                         linter.get_message_state_scope('C0202'))
        linter.disable('W0101', scope='module', line=3)
        self.assertEqual(MSG_STATE_SCOPE_CONFIG,
                         linter.get_message_state_scope('C0202'))
        self.assertEqual(MSG_STATE_SCOPE_MODULE,
                         linter.get_message_state_scope('W0101', 3))
        linter.enable('W0102', scope='module', line=3)
        self.assertEqual(MSG_STATE_SCOPE_MODULE,
                         linter.get_message_state_scope('W0102', 3))
        linter.config = FakeConfig()
        self.assertEqual(
            MSG_STATE_CONFIDENCE,
            linter.get_message_state_scope('this-is-bad',
                                           confidence=interfaces.INFERENCE))

    def test_enable_message_block(self):
        linter = self.init_linter()
        linter.open()
        filepath = join(INPUTDIR, 'func_block_disable_msg.py')
        linter.set_current_module('func_block_disable_msg')
        astroid = linter.get_ast(filepath, 'func_block_disable_msg')
        linter.process_tokens(tokenize_module(astroid))
        fs = linter.file_state
        fs.collect_block_lines(linter.msgs_store, astroid)
        # global (module level)
        self.assertTrue(linter.is_message_enabled('W0613'))
        self.assertTrue(linter.is_message_enabled('E1101'))
        # meth1
        self.assertTrue(linter.is_message_enabled('W0613', 13))
        # meth2
        self.assertFalse(linter.is_message_enabled('W0613', 18))
        # meth3
        self.assertFalse(linter.is_message_enabled('E1101', 24))
        self.assertTrue(linter.is_message_enabled('E1101', 26))
        # meth4
        self.assertFalse(linter.is_message_enabled('E1101', 32))
        self.assertTrue(linter.is_message_enabled('E1101', 36))
        # meth5
        self.assertFalse(linter.is_message_enabled('E1101', 42))
        self.assertFalse(linter.is_message_enabled('E1101', 43))
        self.assertTrue(linter.is_message_enabled('E1101', 46))
        self.assertFalse(linter.is_message_enabled('E1101', 49))
        self.assertFalse(linter.is_message_enabled('E1101', 51))
        # meth6
        self.assertFalse(linter.is_message_enabled('E1101', 57))
        self.assertTrue(linter.is_message_enabled('E1101', 61))
        self.assertFalse(linter.is_message_enabled('E1101', 64))
        self.assertFalse(linter.is_message_enabled('E1101', 66))

        self.assertTrue(linter.is_message_enabled('E0602', 57))
        self.assertTrue(linter.is_message_enabled('E0602', 61))
        self.assertFalse(linter.is_message_enabled('E0602', 62))
        self.assertTrue(linter.is_message_enabled('E0602', 64))
        self.assertTrue(linter.is_message_enabled('E0602', 66))
        # meth7
        self.assertFalse(linter.is_message_enabled('E1101', 70))
        self.assertTrue(linter.is_message_enabled('E1101', 72))
        self.assertTrue(linter.is_message_enabled('E1101', 75))
        self.assertTrue(linter.is_message_enabled('E1101', 77))

        fs = linter.file_state
        self.assertEqual(17, fs._suppression_mapping['W0613', 18])
        self.assertEqual(30, fs._suppression_mapping['E1101', 33])
        self.assertTrue(('E1101', 46) not in fs._suppression_mapping)
        self.assertEqual(1, fs._suppression_mapping['C0302', 18])
        self.assertEqual(1, fs._suppression_mapping['C0302', 50])
        # This is tricky. While the disable in line 106 is disabling
        # both 108 and 110, this is usually not what the user wanted.
        # Therefore, we report the closest previous disable comment.
        self.assertEqual(106, fs._suppression_mapping['E1101', 108])
        self.assertEqual(109, fs._suppression_mapping['E1101', 110])

    def test_enable_by_symbol(self):
        """messages can be controlled by symbolic names.

        The state is consistent across symbols and numbers.
        """
        linter = self.init_linter()
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('unreachable'))
        self.assertTrue(linter.is_message_enabled('W0102'))
        self.assertTrue(linter.is_message_enabled('dangerous-default-value'))
        linter.disable('unreachable', scope='package')
        linter.disable('dangerous-default-value', scope='module', line=1)
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertFalse(linter.is_message_enabled('unreachable'))
        self.assertFalse(linter.is_message_enabled('W0102', 1))
        self.assertFalse(linter.is_message_enabled('dangerous-default-value', 1))
        linter.set_current_module('tutu')
        self.assertFalse(linter.is_message_enabled('W0101'))
        self.assertFalse(linter.is_message_enabled('unreachable'))
        self.assertTrue(linter.is_message_enabled('W0102'))
        self.assertTrue(linter.is_message_enabled('dangerous-default-value'))
        linter.enable('unreachable', scope='package')
        linter.enable('dangerous-default-value', scope='module', line=1)
        self.assertTrue(linter.is_message_enabled('W0101'))
        self.assertTrue(linter.is_message_enabled('unreachable'))
        self.assertTrue(linter.is_message_enabled('W0102', 1))
        self.assertTrue(linter.is_message_enabled('dangerous-default-value', 1))

    def test_lint_ext_module_with_file_output(self):
        self.linter.set_reporter(text.TextReporter())
        if sys.version_info < (3, 0):
            strio = 'StringIO'
        else:
            strio = 'io'
        self.linter.config.files_output = True
        pylint_strio = 'pylint_%s.txt' % strio
        files = [pylint_strio, 'pylint_global.txt']
        for file in files:
            self.addCleanup(remove, file)

        self.linter.check(strio)
        self.linter.generate_reports()
        for f in files:
            self.assertTrue(os.path.exists(f))

    def test_enable_report(self):
        self.assertEqual(self.linter.report_is_enabled('RP0001'), True)
        self.linter.disable('RP0001')
        self.assertEqual(self.linter.report_is_enabled('RP0001'), False)
        self.linter.enable('RP0001')
        self.assertEqual(self.linter.report_is_enabled('RP0001'), True)

    def test_report_output_format_aliased(self):
        text.register(self.linter)
        self.linter.set_option('output-format', 'text')
        self.assertEqual(self.linter.reporter.__class__.__name__, 'TextReporter')

    def test_report_output_format_custom(self):
        this_module = sys.modules[__name__]
        class TestReporter(object):
            pass
        this_module.TestReporter = TestReporter
        class_name = ".".join((this_module.__name__, 'TestReporter'))
        self.linter.set_option('output-format', class_name)
        self.assertEqual(self.linter.reporter.__class__.__name__, 'TestReporter')

    def test_set_option_1(self):
        linter = self.linter
        linter.set_option('disable', 'C0111,W0234')
        self.assertFalse(linter.is_message_enabled('C0111'))
        self.assertFalse(linter.is_message_enabled('W0234'))
        self.assertTrue(linter.is_message_enabled('W0113'))
        self.assertFalse(linter.is_message_enabled('missing-docstring'))
        self.assertFalse(linter.is_message_enabled('non-iterator-returned'))

    def test_set_option_2(self):
        linter = self.linter
        linter.set_option('disable', ('C0111', 'W0234') )
        self.assertFalse(linter.is_message_enabled('C0111'))
        self.assertFalse(linter.is_message_enabled('W0234'))
        self.assertTrue(linter.is_message_enabled('W0113'))
        self.assertFalse(linter.is_message_enabled('missing-docstring'))
        self.assertFalse(linter.is_message_enabled('non-iterator-returned'))

    def test_enable_checkers(self):
        self.linter.disable('design')
        self.assertFalse('design' in [c.name for c in self.linter.prepare_checkers()])
        self.linter.enable('design')
        self.assertTrue('design' in [c.name for c in self.linter.prepare_checkers()])

    def test_errors_only(self):
        linter = self.linter
        self.linter.error_mode()
        checkers = self.linter.prepare_checkers()
        checker_names = set(c.name for c in checkers)
        should_not = set(('design', 'format', 'metrics',
                      'miscellaneous', 'similarities'))
        self.assertSetEqual(set(), should_not & checker_names)

    def test_disable_similar(self):
        self.linter.set_option('disable', 'RP0801')
        self.linter.set_option('disable', 'R0801')
        self.assertFalse('similarities' in [c.name for c in self.linter.prepare_checkers()])

    def test_disable_alot(self):
        """check that we disabled a lot of checkers"""
        self.linter.set_option('reports', False)
        self.linter.set_option('disable', 'R,C,W')
        checker_names = [c.name for c in self.linter.prepare_checkers()]
        for cname in  ('design', 'metrics', 'similarities'):
            self.assertFalse(cname in checker_names, cname)

    def test_addmessage(self):
        self.linter.set_reporter(TestReporter())
        self.linter.open()
        self.linter.set_current_module('0123')
        self.linter.add_message('C0301', line=1, args=(1, 2))
        self.linter.add_message('line-too-long', line=2, args=(3, 4))
        self.assertEqual(
            ['C:  1: Line too long (1/2)', 'C:  2: Line too long (3/4)'],
            self.linter.reporter.messages)

    def test_init_hooks_called_before_load_plugins(self):
        self.assertRaises(RuntimeError,
                          Run, ['--load-plugins', 'unexistant', '--init-hook', 'raise RuntimeError'])
        self.assertRaises(RuntimeError,
                          Run, ['--init-hook', 'raise RuntimeError', '--load-plugins', 'unexistant'])

    def test_analyze_explicit_script(self):
        self.linter.set_reporter(TestReporter())
        self.linter.check(os.path.join(os.path.dirname(__file__), 'data', 'ascript'))
        self.assertEqual(
            ['C:  2: Line too long (175/100)'],
            self.linter.reporter.messages)

    def test_html_reporter_missing_files(self):
        output = six.StringIO()
        with catch_warnings():
            self.linter.set_reporter(html.HTMLReporter(output))

        self.linter.set_option('output-format', 'html')
        self.linter.check('troppoptop.py')
        self.linter.generate_reports()
        value = output.getvalue()
        self.assertIn('troppoptop.py', value)
        self.assertIn('fatal', value)

    def test_python3_checker_disabled(self):
        checker_names = [c.name for c in self.linter.prepare_checkers()]
        self.assertNotIn('python3', checker_names)

        self.linter.set_option('enable', 'python3')
        checker_names = [c.name for c in self.linter.prepare_checkers()]
        self.assertIn('python3', checker_names)

class ConfigTC(unittest.TestCase):

    def setUp(self):
        os.environ.pop('PYLINTRC', None)

    def test_pylint_home(self):
        uhome = os.path.expanduser('~')
        if uhome == '~':
            expected = '.pylint.d'
        else:
            expected = os.path.join(uhome, '.pylint.d')
        self.assertEqual(config.PYLINT_HOME, expected)

        try:
            pylintd = join(tempfile.gettempdir(), '.pylint.d')
            os.environ['PYLINTHOME'] = pylintd
            try:
                reload_module(config)
                self.assertEqual(config.PYLINT_HOME, pylintd)
            finally:
                try:
                    os.remove(pylintd)
                except:
                    pass
        finally:
            del os.environ['PYLINTHOME']

    def test_pylintrc(self):
        with fake_home():
            try:
                self.assertEqual(config.find_pylintrc(), None)
                os.environ['PYLINTRC'] = join(tempfile.gettempdir(),
                                              '.pylintrc')
                self.assertEqual(config.find_pylintrc(), None)
                os.environ['PYLINTRC'] = '.'
                self.assertEqual(config.find_pylintrc(), None)
            finally:
                reload_module(config)

    def test_pylintrc_parentdir(self):
        with tempdir() as chroot:

            create_files(['a/pylintrc', 'a/b/__init__.py', 'a/b/pylintrc',
                          'a/b/c/__init__.py', 'a/b/c/d/__init__.py',
                          'a/b/c/d/e/.pylintrc'])
            with fake_home():
                self.assertEqual(config.find_pylintrc(), None)
            results = {'a'       : join(chroot, 'a', 'pylintrc'),
                       'a/b'     : join(chroot, 'a', 'b', 'pylintrc'),
                       'a/b/c'   : join(chroot, 'a', 'b', 'pylintrc'),
                       'a/b/c/d' : join(chroot, 'a', 'b', 'pylintrc'),
                       'a/b/c/d/e' : join(chroot, 'a', 'b', 'c', 'd', 'e', '.pylintrc'),
                       }
            for basedir, expected in results.items():
                os.chdir(join(chroot, basedir))
                self.assertEqual(config.find_pylintrc(), expected)

    def test_pylintrc_parentdir_no_package(self):
        with tempdir() as chroot:
            with fake_home():
                create_files(['a/pylintrc', 'a/b/pylintrc', 'a/b/c/d/__init__.py'])
                self.assertEqual(config.find_pylintrc(), None)
                results = {'a'       : join(chroot, 'a', 'pylintrc'),
                           'a/b'     : join(chroot, 'a', 'b', 'pylintrc'),
                           'a/b/c'   : None,
                           'a/b/c/d' : None,
                           }
                for basedir, expected in results.items():
                    os.chdir(join(chroot, basedir))
                    self.assertEqual(config.find_pylintrc(), expected)

class PreprocessOptionsTC(unittest.TestCase):
    def _callback(self, name, value):
        self.args.append((name, value))

    def test_value_equal(self):
        self.args = []
        preprocess_options(['--foo', '--bar=baz', '--qu=ux'],
                           {'foo' : (self._callback, False),
                            'qu' : (self._callback, True)})
        self.assertEqual(
            [('foo', None), ('qu', 'ux')], self.args)

    def test_value_space(self):
        self.args = []
        preprocess_options(['--qu', 'ux'],
                           {'qu' : (self._callback, True)})
        self.assertEqual(
            [('qu', 'ux')], self.args)

    def test_error_missing_expected_value(self):
        self.assertRaises(
            ArgumentPreprocessingError,
            preprocess_options,
            ['--foo', '--bar', '--qu=ux'],
            {'bar' : (None, True)})
        self.assertRaises(
            ArgumentPreprocessingError,
            preprocess_options,
            ['--foo', '--bar'],
            {'bar' : (None, True)})

    def test_error_unexpected_value(self):
        self.assertRaises(
            ArgumentPreprocessingError,
            preprocess_options,
            ['--foo', '--bar=spam', '--qu=ux'],
            {'bar' : (None, False)})

class MessagesStoreTC(unittest.TestCase):
    def setUp(self):
        self.store = MessagesStore()
        class Checker(object):
            name = 'achecker'
            msgs = {
                'W1234': ('message', 'msg-symbol', 'msg description.',
                          {'old_names': [('W0001', 'old-symbol')]}),
                'E1234': ('Duplicate keyword argument %r in %s call',
                          'duplicate-keyword-arg',
                          'Used when a function call passes the same keyword argument multiple times.',
                          {'maxversion': (2, 6)}),
                }
        self.store.register_messages(Checker())

    def _compare_messages(self, desc, msg, checkerref=False):
        self.assertMultiLineEqual(desc, msg.format_help(checkerref=checkerref))

    def test_check_message_id(self):
        self.assertIsInstance(self.store.check_message_id('W1234'),
                              MessageDefinition)
        self.assertRaises(UnknownMessage,
                          self.store.check_message_id, 'YB12')

    def test_message_help(self):
        msg = self.store.check_message_id('W1234')
        self._compare_messages(
            ''':msg-symbol (W1234): *message*
  msg description. This message belongs to the achecker checker.''',
            msg, checkerref=True)
        self._compare_messages(
            ''':msg-symbol (W1234): *message*
  msg description.''',
            msg, checkerref=False)

    def test_message_help_minmax(self):
        # build the message manually to be python version independant
        msg = self.store.check_message_id('E1234')
        self._compare_messages(
            ''':duplicate-keyword-arg (E1234): *Duplicate keyword argument %r in %s call*
  Used when a function call passes the same keyword argument multiple times.
  This message belongs to the achecker checker. It can't be emitted when using
  Python >= 2.6.''',
            msg, checkerref=True)
        self._compare_messages(
            ''':duplicate-keyword-arg (E1234): *Duplicate keyword argument %r in %s call*
  Used when a function call passes the same keyword argument multiple times.
  This message can't be emitted when using Python >= 2.6.''',
            msg, checkerref=False)

    def test_list_messages(self):
        sys.stdout = six.StringIO()
        try:
            self.store.list_messages()
            output = sys.stdout.getvalue()
        finally:
            sys.stdout = sys.__stdout__
        # cursory examination of the output: we're mostly testing it completes
        self.assertIn(':msg-symbol (W1234): *message*', output)

    def test_add_renamed_message(self):
        self.store.add_renamed_message('W1234', 'old-bad-name', 'msg-symbol')
        self.assertEqual('msg-symbol',
                         self.store.check_message_id('W1234').symbol)
        self.assertEqual('msg-symbol',
                         self.store.check_message_id('old-bad-name').symbol)

    def test_renamed_message_register(self):
        self.assertEqual('msg-symbol',
                         self.store.check_message_id('W0001').symbol)
        self.assertEqual('msg-symbol',
                         self.store.check_message_id('old-symbol').symbol)

if __name__ == '__main__':
    unittest.main()


from config import bot, chat_id
from time import sleep
from telebot import types
from plugins.error import Error
from plugins.message import shout
import random
from plugins.error import in_chat, check_private

@in_chat()
def say(m):
    bot.delete_message(m.chat.id, m.message_id)
    try:
        if m.chat.type != "private":
            markup = types.InlineKeyboardMarkup() #Отвечаем, если выхов был из супер чата
            link_bot= types.InlineKeyboardButton(text='Перейти в лс', url='t.me/cat_glav_bot') #Отвечаем, если выхов был из супер чата
            markup.add(link_bot) #Отвечаем, если выхов был из супер чата
            sent=bot.send_message(m.chat.id, "Команда /say работает только в лс бота", reply_markup = markup) #Отвечаем, если выхов был из супер чата
            sleep(10)
            bot.delete_message(m.chat.id,sent.message_id)

        if m.reply_to_message:
            sticker_id = bot.get_file(m.reply_to_message.sticker.file_id)
            bot.send_sticker(chat_id, sticker_id.file_id)
            sent=bot.send_message(m.chat.id, "Стикер успешно отправлен!")#Отвечаем, если команда пришла не из супер чата
            sleep(10)
            bot.delete_message(m.chat.id,sent.message_id)
        else:
            bot.send_message(chat_id,  f"_{random.choice(shout)}:_ *'{m.text[5:]}'* 😱 ", parse_mode="Markdown") #Обработать команду и отправить то, что находится с 5 символа и до...
            sent=bot.send_message(m.chat.id, "Сообщение успешно отправлено!")#Отвечаем, если команда пришла не из супер чата
            sleep(10)
            bot.delete_message(m.chat.id,sent.message_id)
    except Exception:
        Error(m, bot).error()


'''
*****************************************************************************************
*
*        =================================================
*             Pharma Bot Theme (eYRC 2022-23)
*        =================================================
*
*  This script is intended for implementation of Task 4A
*  of Pharma Bot (PB) Theme (eYRC 2022-23).
*
*  Filename:			task_4a.py
*  Created:
*  Last Modified:		02/01/2023
*  Author:				e-Yantra Team
*
*  This software is made available on an "AS IS WHERE IS BASIS".
*  Licensee/end user indemnifies and will keep e-Yantra indemnified from
*  any and all claim(s) that emanate from the use of the Software or
*  breach of the terms of this agreement.
*
*****************************************************************************************
'''

# Team ID:			[ 3004 ]
# Author List:		[ Aryan Bawankar, Advait Dhamorikar ]
# Filename:			task_4a.py
# Functions:		[ place_packages, place_traffic_signals, place_start_end_nodes, place_horizontal_barricade, place_vertical_barricade]
#
####################### IMPORT MODULES #######################
## You are not allowed to make any changes in this section. ##
##############################################################
import numpy as np
import cv2
from zmqRemoteApi import RemoteAPIClient
import zmq
import os
import time
##############################################################

################# ADD UTILITY FUNCTIONS HERE #################

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

def place_packages(medicine_package_details, sim, all_models):
    """
	Purpose:
	---
	This function takes details (colour, shape and shop) of the packages present in
    the arena (using "detect_arena_parameters" function from task_1a.py) and places
    them on the virtual arena. The packages should be inserted only into the
    designated areas in each shop as mentioned in the Task document.

    Functions from Regular API References should be used to set the position of the
    packages.

	Input Arguments:
	---
	`medicine_package_details` :	[ list ]
                                nested list containing details of the medicine packages present.
                                Each element of this list will contain
                                - Shop number as Shop_n
                                - Color of the package as a string
                                - Shape of the package as a string
                                - Centroid co-ordinates of the package

    `sim` : [ object ]
            ZeroMQ RemoteAPI object

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	Returns:

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene

	Example call:
	---
	all_models = place_packages(medicine_package_details, sim, all_models)
	"""
    models_directory = os.getcwd()
    packages_models_directory = os.path.join(models_directory, "package_models")
    arena = sim.getObject('/Arena')
####################### ADD YOUR CODE HERE #########################
    flag1 = 0
    flag2 = 0
    flag3 = 0
    flag4 = 0
    flag5 = 0
    for i in medicine_package_details:
        shop = i[0]
        if i[2] == "Circle":
            shape = "cylinder"
        elif i[2] == "Square":
            shape = "cube"
        elif i[2] == "Triangle":
            shape = "cone"
        # Setting Coordinate
        if shop == "Shop_1":
            if flag1 == 1:
                x = x + 0.09
            else:
                x = -0.9 + 0.044
                flag1 = 1
        elif shop == "Shop_2":
            if flag2 == 1:
                x = x + 0.09
            else:
                x = -0.54 + 0.044
                flag2 = 1
        elif shop == "Shop_3":
            if flag3 == 1:
                x = x + 0.09
            else:
                x = -0.18 + 0.044
                flag3 = 1
        elif shop == "Shop_4":
            if flag4 == 1:
                x = x + 0.09
            else:
                x = 0.18 + 0.044
                flag4 = 1
        elif shop == "Shop_5":
            if flag5 == 1:
                x = x + 0.09
            else:
                x = 0.54 + 0.044
                flag5 = 1

        package = i[1] + "_" + shape
        package_ttm = package + ".ttm"
        # print(shop, package)
        package_ttm = os.path.join(packages_models_directory, package_ttm)
        medicine = sim.loadModel(package_ttm)
        sim.setObjectParent(medicine, arena, False)
        sim.setObjectAlias(medicine, package)
        sim.setObjectPosition(medicine, arena, [x, 0.65, 0.015])
        all_models.append(medicine)

####################################################################
    return all_models

def place_traffic_signals(traffic_signals, sim, all_models):
    """
	Purpose:
	---
	This function takes position of the traffic signals present in
    the arena (using "detect_arena_parameters" function from task_1a.py) and places
    them on the virtual arena. The signal should be inserted at a particular node.

    Functions from Regular API References should be used to set the position of the
    signals.

	Input Arguments:
	---
	`traffic_signals` : [ list ]
			list containing nodes in which traffic signals are present

    `sim` : [ object ]
            ZeroMQ RemoteAPI object

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	Returns:

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	None

	Example call:
	---
	all_models = place_traffic_signals(traffic_signals, sim, all_models)
	"""
    models_directory = os.getcwd()
    traffic_sig_model = os.path.join(models_directory, "signals", "traffic_signal.ttm" )
    arena = sim.getObject('/Arena')
####################### ADD YOUR CODE HERE #########################
    for i in traffic_signals:
        a = i[0]
        b = i[1]
        x = 0
        y = 0
        # setting up X Coordinate
        if a == 'A':
            x = -0.9
        elif a == 'B':
            x = -0.54
        elif a == 'C':
            x = -0.18
        elif a == 'D':
            x = 0.18
        elif a == 'E':
            x = 0.54
        elif a == 'F':
            x = 0.9

        # setting up Y Coordinate
        if b == '1':
            y = 0.9
        elif b == '2':
            y = 0.54
        elif b == '3':
            y = 0.18
        elif b == '4':
            y = -0.18
        elif b == '5':
            y = -0.54
        elif b == '6':
            y = -0.9
        name = "Signal_" + i
        position = [x, y, 0.15588]
        signal = sim.loadModel(traffic_sig_model)
        sim.setObjectParent(signal, arena, False)
        sim.setObjectAlias(signal, name)
        sim.setObjectPosition(signal, arena, position)
        all_models.append(signal)
####################################################################
    return all_models

def place_start_end_nodes(start_node, end_node, sim, all_models):
    """
	Purpose:
	---
	This function takes position of start and end nodes present in
    the arena and places them on the virtual arena.
    The models should be inserted at a particular node.

    Functions from Regular API References should be used to set the position of the
    start and end nodes.

	Input Arguments:
	---
	`start_node` : [ string ]
    `end_node` : [ string ]

    `sim` : [ object ]
            ZeroMQ RemoteAPI object

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	Returns:

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	---
	None

	Example call:
	---
	all_models = place_start_end_nodes(start_node, end_node, sim, all_models)
	"""
    models_directory = os.getcwd()
    start_node_model = os.path.join(models_directory, "signals", "start_node.ttm" )
    end_node_model = os.path.join(models_directory, "signals", "end_node.ttm" )
    arena = sim.getObject('/Arena')
####################### ADD YOUR CODE HERE #########################
    a = start_node[0]
    b = start_node[1]
    x = 0
    y = 0
       # setting up X Coordinate
    if a == 'A':
            x = -0.9
    elif a == 'B':
            x = -0.54
    elif a == 'C':
            x = -0.18
    elif a == 'D':
            x = 0.18
    elif a == 'E':
            x = 0.54
    elif a == 'F':
            x = 0.9

        # setting up Y Coordinate
    if b == '1':
            y = 0.9
    elif b == '2':
            y = 0.54
    elif b == '3':
            y = 0.18
    elif b == '4':
            y = -0.18
    elif b == '5':
            y = -0.54
    elif b == '6':
            y = -0.9
    name = "Start_Node"
    position = [x, y, 0.15588]
    start_node = sim.loadModel(start_node_model)
    sim.setObjectParent(start_node, arena, False)
    sim.setObjectAlias(start_node, name)
    sim.setObjectPosition(start_node, arena, position)
    all_models.append(start_node)

    a = end_node[0]
    b = end_node[1]
    x = 0
    y = 0
    # setting up X Coordinate
    if a == 'A':
        x = -0.9
    elif a == 'B':
        x = -0.54
    elif a == 'C':
        x = -0.18
    elif a == 'D':
        x = 0.18
    elif a == 'E':
        x = 0.54
    elif a == 'F':
        x = 0.9

        # setting up Y Coordinate
    if b == '1':
        y = 0.9
    elif b == '2':
        y = 0.54
    elif b == '3':
        y = 0.18
    elif b == '4':
        y = -0.18
    elif b == '5':
        y = -0.54
    elif b == '6':
        y = -0.9
    name = "End_Node"
    position = [x, y, 0.15588]
    end_node = sim.loadModel(end_node_model)
    sim.setObjectParent(end_node, arena, False)
    sim.setObjectAlias(end_node, name)
    sim.setObjectPosition(end_node, arena, position)
    all_models.append(end_node)
####################################################################
    return all_models

def place_horizontal_barricade(horizontal_roads_under_construction, sim, all_models):
    """
	Purpose:
	---
	This function takes the list of missing horizontal roads present in
    the arena (using "detect_arena_parameters" function from task_1a.py) and places
    horizontal barricades on virtual arena. The barricade should be inserted
    between two nodes as shown in Task document.

    Functions from Regular API References should be used to set the position of the
    horizontal barricades.

	Input Arguments:
	---
	`horizontal_roads_under_construction` : [ list ]
			list containing missing horizontal links

    `sim` : [ object ]
            ZeroMQ RemoteAPI object

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	Returns:

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	---
	None

	Example call:
	---
	all_models = place_horizontal_barricade(horizontal_roads_under_construction, sim, all_models)
	"""
    models_directory = os.getcwd()
    horiz_barricade_model = os.path.join(models_directory, "barricades", "horizontal_barricade.ttm" )
    arena = sim.getObject('/Arena')
####################### ADD YOUR CODE HERE #########################
    for i in horizontal_roads_under_construction:
        nodes = i

        A = nodes[0]
         # setting up X Coordinate
        if A == 'A':
            x = -0.9
        elif A == 'B':
            x = -0.54
        elif A == 'C':
            x = -0.18
        elif A == 'D':
            x = 0.18
        elif A == 'E':
            x = 0.54
        elif A == 'F':
            x = 0.9
        x = x + 0.18

        C = nodes[1]
         # setting up Y Coordinate
        if C == '1':
            y = 0.9
        elif C == '2':
            y = 0.54
        elif C == '3':
            y = 0.18
        elif C == '4':
            y = -0.18
        elif C == '5':
            y = -0.54
        elif C == '6':
            y = -0.9

        positions = [x, y, 0.027]
        name = "Horizontal_missing_road_" + A + C + "_" + nodes[3] + nodes[4]
        h_barricade = sim.loadModel(horiz_barricade_model)
        sim.setObjectParent(h_barricade, arena, False)
        sim.setObjectAlias(h_barricade, name)
        sim.setObjectPosition(h_barricade, arena, positions)
        all_models.append(h_barricade)
####################################################################
    return all_models

def place_vertical_barricade(vertical_roads_under_construction, sim, all_models):
    """
	Purpose:
	---
	This function takes the list of missing vertical roads present in
    the arena (using "detect_arena_parameters" function from task_1a.py) and places
    vertical barricades on virtual arena. The barricade should be inserted
    between two nodes as shown in Task document.

    Functions from Regular API References should be used to set the position of the
    vertical barricades.

	Input Arguments:
	---
	`vertical_roads_under_construction` : [ list ]
			list containing missing vertical links

    `sim` : [ object ]
            ZeroMQ RemoteAPI object

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	Returns:

    `all_models` : [ list ]
            list containing handles of all the models imported into the scene
	---
	None

	Example call:
	---
	all_models = place_vertical_barricade(vertical_roads_under_construction, sim, all_models)
	"""
    models_directory = os.getcwd()
    vert_barricade_model = os.path.join(models_directory, "barricades", "vertical_barricade.ttm" )
    arena = sim.getObject('/Arena')
####################### ADD YOUR CODE HERE #########################
    for i in vertical_roads_under_construction:
        nodes = i

        A = nodes[0]
         # setting up X Coordinate
        if A == 'A':
            x = -0.9
        elif A == 'B':
            x = -0.54
        elif A == 'C':
            x = -0.18
        elif A == 'D':
            x = 0.18
        elif A == 'E':
            x = 0.54
        elif A == 'F':
            x = 0.9

        C = nodes[1]
         # setting up Y Coordinate
        if C == '1':
            y = 0.9
        elif C == '2':
            y = 0.54
        elif C == '3':
            y = 0.18
        elif C == '4':
            y = -0.18
        elif C == '5':
            y = -0.54
        elif C == '6':
            y = -0.9
        y = y - 0.18

        positions = [x, y, 0.027]
        name = "Vertical_missing_road_" + A + C + "_" + nodes[3] + nodes[4]
        v_barricade = sim.loadModel(vert_barricade_model)
        sim.setObjectParent(v_barricade, arena, False)
        sim.setObjectAlias(v_barricade, name)
        sim.setObjectPosition(v_barricade, arena, positions)
        all_models.append(v_barricade)
####################################################################
    return all_models

if __name__ == "__main__":
    client = RemoteAPIClient()
    sim = client.getObject('sim')
    # arena = sim.getObject('/Arena')
    aruco_handle = sim.getObject('/aruco_3')

    arena = sim.getObject('/Arena')
    # sim.setObjectParent(aruco_handle, arena, False)
    # sim.setObjectAlias(aruco_handle, "marker")
    sim.setObjectPosition(aruco_handle, -1, [0.15, 0.15, 0.15])
    sim.setObjectOrientation(aruco_handle, -1, [0, 0, 45])
    # path directory of images in test_images folder
    img_dir = os.getcwd() + "/test_imgs/"

    i = 0
    config_img = cv2.imread(img_dir + 'maze_' + str(i) + '.png')

    print('\n============================================')
    print('\nFor maze_0.png')

    # object handles of each model that gets imported to the scene can be stored in this list
    # at the end of each test image, all the models will be removed
    all_models = []

    # import task_1a.py. Make sure that task_1a.py is in same folder as task_4a.py
    task_1 = __import__('task_1a')
    detected_arena_parameters = task_1.detect_arena_parameters(config_img)

    # obtain required arena parameters
    medicine_package_details = detected_arena_parameters["medicine_packages"]
    traffic_signals = detected_arena_parameters['traffic_signals']
    start_node = detected_arena_parameters['start_node']
    end_node = detected_arena_parameters['end_node']
    horizontal_roads_under_construction = detected_arena_parameters[
        'horizontal_roads_under_construction']
    vertical_roads_under_construction = detected_arena_parameters[
        'vertical_roads_under_construction']

    print("[1] Setting up the scene in CoppeliaSim")
    all_models = place_packages(medicine_package_details, sim, all_models)
    all_models = place_traffic_signals(traffic_signals, sim, all_models)
    all_models = place_horizontal_barricade(
        horizontal_roads_under_construction, sim, all_models)
    all_models = place_vertical_barricade(
        vertical_roads_under_construction, sim, all_models)
    all_models = place_start_end_nodes(start_node, end_node, sim, all_models)
    print("[2] Completed setting up the scene in CoppeliaSim")

    # wait for 10 seconds and then remove models
    time.sleep(10)
    print("[3] Removing models for maze_0.png")

    for i in all_models:
        sim.removeModel(i)

    choice = input(
        '\nDo you want to run your script on all test images ? => "y" or "n": ')

    if choice == 'y':
        for i in range(1, 5):

            print('\n============================================')
            print('\nFor maze_' + str(i) + '.png')
            config_img = cv2.imread(img_dir + 'maze_' + str(i) + '.png')

            # object handles of each model that gets imported to the scene can be stored in this list
            # at the end of each test image, all the models will be removed
            all_models = []

            # import task_1a.py. Make sure that task_1a.py is in same folder as task_4a.py
            task_1 = __import__('task_1a')
            detected_arena_parameters = task_1.detect_arena_parameters(
                config_img)

            # obtain required arena parameters
            medicine_package_details = detected_arena_parameters["medicine_packages"]
            traffic_signals = detected_arena_parameters['traffic_signals']
            start_node = detected_arena_parameters['start_node']
            end_node = detected_arena_parameters['end_node']
            horizontal_roads_under_construction = detected_arena_parameters[
                'horizontal_roads_under_construction']
            vertical_roads_under_construction = detected_arena_parameters[
                'vertical_roads_under_construction']

            print("[1] Setting up the scene in CoppeliaSim")
            place_packages(medicine_package_details, sim, all_models)
            place_traffic_signals(traffic_signals, sim, all_models)
            place_horizontal_barricade(
                horizontal_roads_under_construction, sim, all_models)
            place_vertical_barricade(
                vertical_roads_under_construction, sim, all_models)
            place_start_end_nodes(start_node, end_node, sim, all_models)
            print("[2] Completed setting up the scene in CoppeliaSim")

            # wait for 10 seconds and then remove models
            time.sleep(10)
            print("[3] Removing models for maze_" + str(i) + '.png')
            for i in all_models:
                sim.removeModel(i)


from pydantic import BaseSettings, BaseModel

from pathlib import Path

from .languages import WIKIPEDIA_LANGS

import toml

class Settings(BaseSettings):
    status: str = "unknown"
    logging_chat: int = None

    db_path: Path = Path("jdanbot.db")
    music_path: Path = Path("media/music")

    admin_notes: list[str]

    bot_owners: list[int] = [795449748, 0]

    class Tokens(BaseModel):
        bot_token: str

    class Schedule(BaseModel):
        delay_seconds: int = 20

        katz_bots: bool = False

    class Egg(BaseModel):
        commands: list[str]
        audio: Path

    tokens: Tokens
    schedule: Schedule = Schedule()
    eggs: list[Egg]

with open("settings.toml") as file:
    settings_file = toml.loads(file.read())

with open(".secrets.toml") as file:
    secrets_file = toml.loads(file.read())

settings = Settings.parse_obj(settings_file | secrets_file)

BASE_DIR = Path(__file__).parent.parent.parent
LOCALES_DIR = BASE_DIR / "locales"

WIKIPEDIA_SHORTCUTS = {
    "ru": ["w"],
    "en": ["v"],
    "uk": ["wua", "wikiua"]
}

WIKI_COMMANDS = []

for lang in WIKIPEDIA_LANGS:
    WIKI_COMMANDS.extend([f"wiki{lang}", f"w{lang}"])

for lang in WIKIPEDIA_SHORTCUTS:
    WIKI_COMMANDS.extend(WIKIPEDIA_SHORTCUTS[lang])


from extras import *
from common import Common
from indexer import Indexer
import math
import re
import nltk
import operator
from collections import defaultdict, OrderedDict
from nltk.stem import SnowballStemmer
from nltk.corpus import wordnet

class Query_Expansion:

    def __init__(self):
        """
        Constructor: Used to initialize all the class variables
        """
        self.utility = Utility()
        self.frequency_map = defaultdict()
        self.synonyms_map = defaultdict()
        self.file_handling = FileHandling()
        self.common = Common()
        self.indexer = Indexer()

    def generate_expected_words_for_expansion(self, queries):
        stopWords = self.utility.get_stop_list()
        stemmer = SnowballStemmer("english")
        for i in range (0,len(queries)):
            query = queries[i]
            listofwords = []
            words = query.split()
            for word in words:
                word = word.lower()
                stem = stemmer.stem(word)
                expected = self.fetch_expected_words(word,stem)
                if expected not in stopWords:
                    frequency = self.generate_frequency_map(word,expected)
                    if frequency > 0:
                        listofwords.append(expected)

            self.frequency_map[i+1] = listofwords
        return self.frequency_map

    def generate_frequency_map(self,word,stem):
        occurrences = 0
        if stem in self.positional_index and word in self.positional_index:
            dict_stem = self.positional_index[stem]
            dict_word = self.positional_index[word]

            for doc in dict_word:
                if doc in dict_stem:
                    list1 = dict_word[doc]
                    list2 = dict_stem[doc]
                    pos1 = 0
                    for i in range(0, len(list1)):
                        pos1 = pos1 + list1[i]
                        pos2 = 0
                        for j in range(0, len(list2)):
                            pos2 = pos2 + list2[j]
                            if abs(pos1 - pos2) <= 12:
                                occurrences = occurrences + 1
                                break
        return occurrences

    def fetch_expected_words(self,word,stem):
        if self.utility.check_word_exist(stem):
            return stem
        else:
            return nltk.stem.WordNetLemmatizer().lemmatize(word)

    def expand_queries_using_stemming(self, queries):
        self.positional_index  = self.indexer.read_index(index_type=True)
        print('\n' + self.utility.line_break + '\n' +\
            'Running Query Expansion using Stemming..')
        stem_map = self.generate_expected_words_for_expansion(queries)
        updated_query_map = defaultdict(set)
        for i in range(len(queries)):
            stop_words = self.utility.get_stop_list()
            listofwords = stem_map[i+1]
            for word in listofwords:
                for syn in wordnet.synsets(word):
                    for l in syn.lemmas():
                        if str(l.name) not in  queries[i] and '_' not in str(l.name) and str(l.name) not in stop_words:
                            updated_query_map[i+1].add(l.name())
                            if (len(updated_query_map[i+1])) > 4:
                                break
                    if len(updated_query_map[i+1]) > 4:
                        break

        new_queries = []

        for i in range (len(queries)):
            old_query = queries[i]
            new_query = old_query
            for word in updated_query_map[i+1]:
                new_query = new_query + " "+ str(word)
            new_queries.append(new_query)
        return new_queries

    def create_tf(self,inverted_index):
        tf = {}
        for term in inverted_index:
            c = 0
            doc_to_frequency = inverted_index[term]
            for doc in doc_to_frequency:
                c = c + doc_to_frequency[doc]
            tf[term] = c
        return self.generatePotentialQuery(tf)

    # generating potential query words by evaluating term frequency and removing stop words
    def generatePotentialQuery(self,tf):
        terms = []
        total = 0
        for key, value in tf.items():
            total = total + value
        potentialList = []
        for key, value in tf.items():
            if key not in self.utility.get_stop_list() and len(key) > 4:
                potentialList.append(key)
        return potentialList

    # calculating dice's co-efficient for different terms
    def diceCoff(self,list1, list2, invertedIndex):
        associationDict = {}
        for i in list1:
            if i != "in" and i in invertedIndex:
                docList = invertedIndex[i]
                sum = 0
                for j in list2:
                    docList2 = invertedIndex[j]
                    sum = 0
                    for k in docList2:
                        if k in docList:
                            sum = sum + 1
                    if sum > 10:
                        associationDict[i + "   " + j] = sum * 1.0 / (len(docList) + len(docList2))
        sorted_dict = OrderedDict(associationDict)
        return sorted_dict

    def expand_queries_using_pseduo_relevance(self, queries):
        print('\n' + self.utility.line_break + '\n' +\
            'Running Query Expansion using Pseduo Relevance..')
        docs = self.common.read_top_documents_for_score(top=40)
        relevant_docs = []
        for record in docs:
            relevant_docs.append((record.values()[0]))

        self.indexer.create_save_indexer_with_relevant_docs(relevant_docs)
        inverted_index = self.indexer.read_simple_index()

        potential_list = self.create_tf(inverted_index)
        updated_query_list = []

        for i in range(len(queries)):
            query = queries[i]
            query = query.lower()
            words_from_query = []
            word_array = query.split()
            for word in word_array:
                word = re.sub(r'\W+', ' ', word)
                if word not in self.utility.get_stop_list():
                    words_from_query.append(word)
            updatedQuery = query
            suggested_words = self.diceCoff(words_from_query,potential_list,inverted_index).items()
            k = 0
            for value in suggested_words:
                if k > 8:
                    break
                else:
                    words = value[0].split()
                    if words[1] not in updatedQuery:
                        updatedQuery = updatedQuery + ' ' + words[1]
                        k = k + 1
            updated_query_list.append(updatedQuery)
        return updated_query_list


from sklearn import datasets
from sklearn.impute import SimpleImputer
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np

# Multiclass Classification Datasets

def load_mnist():
    mnist = datasets.load_digits(as_frame=True)
    mnist_X, mnist_y = mnist.data, mnist.target

    X_train, X_test, y_train, y_test = train_test_split(mnist_X, mnist_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_forest_covertypes():
    forest = datasets.fetch_covtype(as_frame=True)
    forest_X, forest_y = forest.data, forest.target
    forest_X = forest_X[:15000]
    forest_y = forest_y[:15000]

    X_train, X_test, y_train, y_test = train_test_split(forest_X, forest_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_kepler_exoplanets():
    kepler = pd.read_csv('data/kepler_exoplanet.csv')
    not_X_columns = ['rowid', 'kepid', 'kepoi_name', 'kepler_name', 'koi_score', 'koi_pdisposition', 'koi_disposition', 'koi_teq_err1', 'koi_teq_err2', 'koi_tce_delivname']

    kepler_X = kepler.drop(not_X_columns, axis=1)
    kepler_y = kepler['koi_pdisposition']

    imputer = SimpleImputer(missing_values=np.nan, strategy='most_frequent')
    kepler_X_array = imputer.fit_transform(kepler_X)
    kepler_X = pd.DataFrame(kepler_X_array, columns=kepler_X.columns)

    X_train, X_test, y_train, y_test = train_test_split(kepler_X, kepler_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

# Regression Datasets

def load_cali_housing():
    california = datasets.fetch_california_housing(as_frame=True)
    cali_X, cali_y = california.data, california.target

    X_train, X_test, y_train, y_test = train_test_split(cali_X, cali_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_melbourne_housing():
    melbourne = pd.read_csv('data/melbourne_housing_data.csv')

    not_X_columns = ['Address', 'Price', 'SellerG', 'Date']
    melbourne_X = melbourne.drop(not_X_columns, axis=1)
    melbourne_y = melbourne['Price']

    # Convert all features to integers
    categorical_cols = ['Suburb', 'Type', 'Method', 'CouncilArea', 'Regionname']
    for col_name in categorical_cols:
        melbourne_X[col_name] = pd.Categorical(melbourne_X[col_name]).codes

    # Impute missing values
    imputer = SimpleImputer(missing_values=np.nan, strategy='most_frequent')
    melbourne_X_array = imputer.fit_transform(melbourne_X)
    melbourne_X = pd.DataFrame(melbourne_X_array, columns=melbourne_X.columns)

    X_train, X_test, y_train, y_test = train_test_split(melbourne_X, melbourne_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_world_happiness():
    happiness = pd.read_csv('data/world-happiness-report-2021.csv')

    X_columns = ['Regional indicator', 'Logged GDP per capita', 'Social support', 'Healthy life expectancy', 'Freedom to make life choices', 'Generosity', 'Perceptions of corruption']
    happiness_X = happiness[X_columns]
    happiness_y = happiness['Ladder score']

    happiness_X['Regional indicator'] = pd.Categorical(happiness_X['Regional indicator']).codes

    X_train, X_test, y_train, y_test = train_test_split(happiness_X, happiness_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

# Binary Classification Datasets

def load_heart_attack():
    heart_attack = pd.read_csv('data/heart-attack.csv')

    heart_X = heart_attack.drop('output', axis=1)
    heart_y = heart_attack['output']

    X_train, X_test, y_train, y_test = train_test_split(heart_X, heart_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_stroke():
    stroke = pd.read_csv('data/healthcare-dataset-stroke-data.csv')

    not_X_columns = ['id', 'stroke']
    stroke_X = stroke.drop(not_X_columns, axis=1)
    stroke_y = stroke['stroke']

    # Convert all features to integers
    categorical_cols = ['gender', 'ever_married', 'work_type', 'Residence_type', 'smoking_status']
    for col_name in categorical_cols:
        stroke_X[col_name] = pd.Categorical(stroke_X[col_name]).codes

    # Impute missing values
    imputer = SimpleImputer(missing_values=np.nan, strategy='most_frequent')
    stroke_X_array = imputer.fit_transform(stroke_X)
    stroke_X = pd.DataFrame(stroke_X_array, columns=stroke_X.columns)

    X_train, X_test, y_train, y_test = train_test_split(stroke_X, stroke_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_telecom():
    telecom = pd.read_csv('data/telecom_users.csv')

    not_X_columns = ['Unnamed: 0', 'customerID', 'TotalCharges', 'Churn']
    telecom_X = telecom.drop(not_X_columns, axis=1)
    telecom_y = pd.Categorical(telecom['Churn']).codes

    # Convert all features to integers
    not_categorical_cols = ['SeniorCitizen', 'tenure', 'MonthlyCharges']
    for col_name in telecom_X.columns:
        if col_name not in not_categorical_cols:
            telecom_X[col_name] = pd.Categorical(telecom_X[col_name]).codes

    # Impute missing values
    imputer = SimpleImputer(missing_values=np.nan, strategy='most_frequent')
    telecom_X_array = imputer.fit_transform(telecom_X)
    telecom_X = pd.DataFrame(telecom_X_array, columns=telecom_X.columns)

    X_train, X_test, y_train, y_test = train_test_split(telecom_X, telecom_y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

def load_multiclass():
    all_data = {}

    all_data['mnist'] = load_mnist()
    all_data['forest_covertypes'] = load_forest_covertypes()
    all_data['kepler_exoplanets'] = load_kepler_exoplanets()

    return all_data

def load_regression():
    all_data = {}

    all_data['california_housing'] = load_cali_housing()
    all_data['melbourne_housing'] = load_melbourne_housing()
    all_data['world_happiness'] = load_world_happiness()

    return all_data

def load_binary():
    all_data = {}

    all_data['heart_attack'] = load_heart_attack()
    all_data['stroke'] = load_stroke()
    all_data['telecom'] = load_telecom()

    return all_data


# pip install uszipcode
import pandas as pd
from uszipcode import SearchEngine
search = SearchEngine(simple_zipcode=False)

def add_coordinates(df):
    '''Input: a pandas dataframe which includes a 'Zip' field representing a
       US zip code.
       Output: dataframe with 'lat' and 'lng' fields added which represent
       latitude and longitude

    '''
    df['lat'] = df['Zip'].apply(lambda x: pd.Series(search.by_zipcode(x).lat))
    df['lng'] = df['Zip'].apply(lambda x: pd.Series(search.by_zipcode(x).lng))

    df['county'] = df['Zip'].apply(lambda x: pd.Series(search.by_zipcode(x).county))

def add_fips(df):
    '''Input: a pandas dataframe which includes a 'Zip' field representing a
       US zip code.
       Output: dataframe with 'lat' and 'lng' fields added which represent
       latitude and longitude

    '''

    colorado_fips = ['08001', '08003', '08005', '08007', '08009', '08011', '08013', '08014', '08015', '08017', '08019',
                    '08021', '08023', '08025', '08027', '08029', '08031', '08033', '08035', '08037', '08039', '08041',
                    '08043', '08045', '08047', '08049', '08051', '08053', '08055', '08057', '08059', '08061', '08063',
                    '08065', '08067', '08069', '08071', '08073', '08075', '08077', '08079', '08081', '08083', '08085',
                    '08087', '08089', '08091', '08093', '08095', '08097', '08099', '08101', '08103', '08105', '08107',
                    '08109', '08111', '08113', '08115', '08117', '08119', '08121', '08123', '08125']

    colorado_counties = ['Adams County', 'Alamosa County', 'Arapahoe County', 'Archuleta County', 'Baca County', 'Bent County', 'Boulder County', 'Broomfield County',
                         'Chaffee County', 'Cheyenne County', 'Clear Creek County', 'Conejos County', 'Costilla County', 'Crowley County', 'Custer County',
                         'Delta County', 'Denver County', 'Dolores County', 'Douglas County', 'Eagle County', 'Elbert County', 'El Paso County', 'Fremont County',
                         'Garfield County', 'Gilpin County', 'Grand County', 'Gunnison County', 'Hinsdale County', 'Huerfano County', 'Jackson County', 'Jefferson County',
                         'Kiowa County', 'Kit Carson County', 'Lake County', 'La Plata County', 'Larimer County', 'Las Animas County',  'Lincoln County', 'Logan County',
                         'Mesa County', 'Mineral County', 'Moffat County', 'Montezuma County', 'Montrose County', 'Morgan County', 'Otero County', 'Ouray County',
                         'Park County', 'Phillips County', 'Pitkin County', 'Prowers County',  'Pueblo County', 'Rio Blanco County', 'Rio Grande County','Routt County',
                         'Saguache County', 'San Juan County', 'San Miguel County', 'Sedgwick County', 'Summit County', 'Teller County', 'Washington County',
                         'Weld County', 'Yuma County']

    counties_fips_dict = dict(zip(colorado_counties, colorado_fips))

    df['county'] = df['Zip'].apply(lambda x: pd.Series(search.by_zipcode(x).county))
    col_only_df = df[df['county'].isin(colorado_counties)]
    col_only_df['fip'] = col_only_df['county'].apply(lambda x: counties_fips_dict[x])
    return col_only_df


import pygame
import Snake, Apple
import os

os.environ['SDL_VIDEO_CENTERED'] = '1'

pygame.mixer.pre_init(44100, -16, 1, 512)
pygame.init()
pygame.display.set_caption("Sssssnake")

# 10x10 segments
winwidth = 200
winheight = 240
win = pygame.display.set_mode((winwidth, winheight))

# segment - 20x20px
segsize = 20
snake = Snake.Snake(0, 20, segsize, segsize-2, segsize-2)
apple = Apple.Apple(segsize//2, winwidth, winheight, segsize, snake.segments)

# font
font = pygame.font.SysFont("monospace", 15)

# sounds
eatsound = pygame.mixer.Sound('sounds/eat sound.wav')
losesound = pygame.mixer.Sound('sounds/lose sound.wav')
music = pygame.mixer.music.load('sounds/bg music.mp3')
pygame.mixer.music.play(-1)

def lost():
    pygame.mixer.music.stop()
    losesound.play()
    global win, running, snake, score
    gameover = font.render("GAME OVER :(", 1, (255, 255, 255))
    playagain = font.render("Play again?", 1, (255, 255, 255))
    yorn = font.render("(y)   (n)", 1, (255, 255, 255))
    win.blit(gameover, (winwidth//2 - 55, winheight//2 - 35))
    win.blit(playagain, (winwidth//2 - 52, winheight//2 - 15))
    win.blit(yorn, (winwidth//2 - 40, winheight//2))
    pygame.display.update()
    pygame.event.clear()
    while True:
        event = pygame.event.wait()
        if event.type == pygame.QUIT:
            running = False
            break
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_y:
                snake = Snake.Snake(0, 20, segsize, segsize-2, segsize-2)
                apple.forbidden = snake.segments
                score = 0
                pygame.mixer.music.play(-1)
                break
            if event.key == pygame.K_n:
                running = False
                break

# mainloop
running = True
score = 0
while running:
    pygame.time.delay(200)
    keypressed = False
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN and not keypressed:
            if event.key == pygame.K_UP and snake.direction != (1, 0):
                snake.direction = (-1, 0)
                keypressed = True
            elif event.key == pygame.K_DOWN and snake.direction != (-1, 0):
                snake.direction = (1, 0)
                keypressed = True
            elif event.key == pygame.K_LEFT and snake.direction != (0, 1):
                snake.direction = (0, -1)
                keypressed = True
            elif event.key == pygame.K_RIGHT and snake.direction != (0, -1):
                snake.direction = (0, 1)
                keypressed = True

    # calculating new position
    tempx = snake.x + snake.vel * snake.direction[1]
    tempy = snake.y + snake.vel * snake.direction[0]
    if 0 <= tempx <= winwidth - snake.segwidth and segsize <= tempy <= winheight - segsize - snake.segheigth:
        snake.x = tempx
        snake.y = tempy
    else:
        #collision with borders
        lost()
        continue

    snake.move()

    # collision with snake
    if snake.segments.count((snake.x, snake.y)) > 1:
        lost()
        continue

    win.fill((0, 0, 0))
    # collision with apple
    if snake.x <= apple.x <= snake.x + segsize and snake.y <= apple.y <= snake.y + segsize:
        eatsound.play()
        snake.addsegment()
        snake.draw(win)
        apple.setposition()
        score += 1
    else:
        snake.draw(win)

    label = font.render("Score {}".format(score), 1, (255, 255, 255))
    win.blit(label, (winwidth - 70, 0))
    apple.draw(win)
    pygame.draw.rect(win, (255, 255, 0), (1, 20, winwidth - 2, winheight - 2*segsize), 1)
    pygame.display.update()

pygame.quit()


import pygame
import jumpingHorses.constants as Constants
from .constants import BLACK_PIECE, WHITE_PIECE, MOVE_COLOR, MOVE_RADIUS, SQUARE_SIZE, WIDTH, HEIGHT, LETTER_GAP_SIZE, OUTLINE_SIZE
from .board import Board
from menu.main_menu import MainMenu
import menu.gameState as GameState

class GameMaster:
    #inicializē
    def __init__(self, surface):
        self._init()
        self.surface = surface
    #atjaunina to, kas redzams ekrānā
    def update(self):
        self.board.draw(self.surface)
        if self.selectedPiece != None:
            self.draw_valid_moves(self.valid_moves)
    #sāk no jauna spēli
    def reset(self):
        self._init()
    #notīra esošo stāvokli
    def _init(self):
        self.selectedPiece = None
        self.turn = WHITE_PIECE
        self.board = Board()
        self.valid_moves = {}
    #pārbauda, vai var izvēlēties
    def select(self, pos):
        if LETTER_GAP_SIZE < pos[0] < WIDTH+LETTER_GAP_SIZE and pos[1] < HEIGHT:
            row, col = self.get_row_col_from_mouse(pos)
            if self.selectedPiece:#ja kaut kas jau ir izvēlēts
                result = self._move(row, col)#tad to pabīda, ja ir legāls gājiens
                self.selectedPiece = None
                if not result:#ja neidzodas pabīdīt(izvēlas neiespējamu gājienu, vai kaut ko citu), tad selecto pa jaunam
                    self.selectedPiece = None
                    self.select(pos)
                    return True
            #ja kaut kas tiek izvēlēts pirmo reizi, vai tika izvēlēts kaut kas, kas nebija iespējams gājiens
            piece = self.board.get_piece(row, col)
            if piece != 0 and piece.color == self.turn:
                self.selectedPiece = piece
                self.valid_moves = self.board.get_valid_moves(piece)
                return True
        self.selectedPiece = None
        return False
    #pārvieto kauliņu
    def _move (self, row, col):
        piece = self.board.get_piece(row,col)
        if self.selectedPiece and piece == 0 and (row,col) in self.valid_moves:
            self.board.move(self.selectedPiece, row, col)
            self.change_turn()
        else:
            return False

        return True
    #samaina gājienu
    def change_turn(self):
        self.valid_moves = {}
        if self.turn == BLACK_PIECE:
            self.turn = WHITE_PIECE
        else:
            self.turn = BLACK_PIECE

    #uzzīmē legālos gājienus
    def draw_valid_moves(self, moves):
        pygame.draw.circle(self.surface, (0,255,0), (self.selectedPiece.x, self.selectedPiece.y), self.selectedPiece.radius+OUTLINE_SIZE, 5)
        for move in moves:
            row, col = move
            pygame.draw.circle(self.surface, MOVE_COLOR, (col*SQUARE_SIZE + SQUARE_SIZE//2, row*SQUARE_SIZE + SQUARE_SIZE//2) , MOVE_RADIUS)

    #atgriež rindu un kolonnu atkarībā no peles pozīcijas
    def get_row_col_from_mouse(self, pos):
        x, y = pos
        row = y // SQUARE_SIZE
        col = (x - LETTER_GAP_SIZE) // SQUARE_SIZE
        return row, col
    #dators veic gājienu
    def ai_move(self, board):
        self.board = board
        self.change_turn()
    #pārbauda, vai ir uzvarētājs
    def check_winner(self):
        if self.board.winner() == Constants.starting_player[1]:
            GameState.currentState = GameState.State.win
        elif self.board.winner() == Constants.starting_player[2]:
            GameState.currentState = GameState.State.lost

    def get_board(self):
        #print("get_board called")
        #self.board.print_board()
        return self.board