from helpers.verifications import validate_content, title_in_blacklist
from bs4 import BeautifulSoup
from datetime import datetime
import requests

def validate_date_bitcoinist(html):
    # Encuentra el div con la clase 'jeg_meta_date'
    date_div = html.find('div', class_='jeg_meta_date')

    if date_div:
        # Encuentra la etiqueta 'a' dentro del div
        date_link = date_div.find('a')
        if date_link:
            date_text = date_link.text.lower()

            # Comprueba si el texto contiene "hours ago"
            if "hours ago" in date_text:
                return date_text

    return False

def extract_image_urls(html):
    image_urls = []
    soup = BeautifulSoup(html, 'html.parser')
    img_elements = soup.find_all('img')

    for img in img_elements:
        src = img.get('src')

        if src and src.startswith('https://bitcoinist.com/wp-content/uploads/'):
            image_urls.append(src)

    return image_urls

# Function to validate the article using keywords
def validate_bitcoinist_article(article_link, main_keyword):

    headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.61 Safari/537.36'
        }

    article_response = requests.get(article_link, headers=headers)
    article_content_type = article_response.headers.get("Content-Type", "").lower()

    if article_response.status_code == 200 and 'text/html' in article_content_type:
        article_soup = BeautifulSoup(article_response.text, 'html.parser')

        title_element = article_soup.find('h1')
        title = title_element.text.strip() if title_element else None

        content = ""
        all_p_elements = article_soup.findAll("p")
        for el in all_p_elements:
            content += el.text.lower()

        if not title or not content:
            # print('Article does not have a title or content')
            return None, None, None, None
        else:
            is_title_in_blacklist = title_in_blacklist(title)
            content_validation = validate_content(main_keyword, content)

        if is_title_in_blacklist or not content_validation:
            # print('Article does not meet requirements')
            return None, None, None, None

        valid_date = validate_date_bitcoinist(article_soup)

        # Extract image URLs from the article
        image_urls = extract_image_urls(article_response.text)

        if  content_validation and valid_date and title:
                return title, content, valid_date, image_urls
        else:
            return None, None, None, None


from typing import List

def max_element(l: List[int]) -> int:
    """Return maximum element in the list.
    >>> max_element([1, 2, 3])
    3
    >>> max_element([5, 3, -5, 2, -3, 3, 9, 0, 123, 1, -10])
    123
    """
    if not l:
        raise ValueError("Cannot find maximum element in empty list.")
    max_so_far = l[0]
    for num in l:
        if num > max_so_far:
            max_so_far = num
    return max_so_far

def check(candidate):
    assert candidate([1, 2, 3]) == 3
    assert candidate([5, 3, -5, 2, -3, 3, 9, 0, 124, 1, -10]) == 124

def test_check():
    check(max_element)

test_check()


"""
File_id: 07oct2021_async_server
Related file id:  07oct2021_async_client, 07oct2021_menu
This is the alpha server code for the "shooter" game
"""

import asyncio
import json
from zlib import compress, decompress
import pygame
import coloredlogs
import logging
from logging import handlers
from dataclasses import dataclass, asdict, field
from typing import Any
from time import perf_counter

from inspect import currentframe, getframeinfo

import game_class_s
import game_state
from platform_shooter_settings import *
# import platform_shooter_settings

FRAME_INFO = getframeinfo(currentframe())

@dataclass
class RoomState:
    # player0: player[0], player1: player[1]
    player_names: list[str] = field(default_factory=list)  # name:str, room name will be f"{player_name[0]}'s game"
    # player_role_ids: list[int] = field(default_factory=list)  # [player0, player1]
    player_role_ids = [0, 1]  # [player0, player1], 0: shooter, 1: chopper
    player_readers: list[asyncio.streams.StreamReader] = field(default_factory=list)  # [player0, player1]
    player_writers: list[asyncio.streams.StreamWriter] = field(default_factory=list)  # [player0, player1]
    player_task_names: list[str] = field(default_factory=list)
    room_id: int = 0
    player_joined: bool = False  # True if the required info is received from 2nd player
    game_set0: bool = False  # True if player0 finished setting map, match, role
    game_set1: bool = False  # True if player 1 finished setting role
    map_id: int = 0
    match_id: int = 0
    running: bool = False
    winner: str = "none"  # the winner name of the round
    choosable: bool = True

    def check_ready(self):
        return self.player_joined and self.game_set0

class Server:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), flags=pygame.HIDDEN)
        self.clock = pygame.time.Clock()
        self.cnt = 0  # total number of connections to the server
        self.client_id = 0
        self.game_dict = {}  # used to store game room information
        self.room_cnt = 0  # total number of game rooms
        self.room_id = 0
        self.player_info = []  # f"{conn_type},{player_name}"
        self.writer: [asyncio.streams.StreamWriter] = None
        self.reader: [asyncio.streams.StreamReader] = None
        self.loop = None  # asyncio loop
        self.my_logger = logging.getLogger()
        fmt = "%(asctime)s [%(levelname)s]: %(message)s"
        coloredlogs.install(level=logging.INFO,
                            logger=self.my_logger,
                            fmt=fmt,
                            field_styles=FIELD_STYLES,
                            level_styles=LEVEL_STYLES)
        format_str = logging.Formatter(fmt)
        fh = handlers.RotatingFileHandler("server_log.txt", "a", 100000, 3)
        fh.setFormatter(format_str)
        # self.my_logger.addHandler(fh)
        self.timer = 0

    async def check_read(self, player_name, reader, writer, in_room_flg):
        # used for read() before the connected player is in a game room
        string = None
        connected = CONNECTED
        try:
            # received = await reader.read(length)
            received = await reader.readuntil(separator=b"AB")
            string = received.decode()[:-2]
        except (ConnectionError, asyncio.IncompleteReadError):
            if not in_room_flg:
                self.cnt -= 1
            self.my_logger.warning(
                f"Connection to player {player_name} is lost [{getframeinfo(currentframe()).lineno}]")
            connected = not CONNECTED
        else:  # if the connection is properly connected, there will be no ConnectionError exception raised
            if not received:
                if not in_room_flg:
                    self.cnt -= 1
                self.my_logger.warning(
                    f"Connection to player {player_name} is lost [{getframeinfo(currentframe()).lineno}]")
                if not writer.is_closing():
                    writer.close()
                    await writer.wait_closed()
                connected = not CONNECTED
        return connected, string

    def remove_room(self, room):
        self.cnt -= 2  # if a room is removed, two players are removed
        # room.running = False
        # self.my_logger.warning(
        #     f"Connection to player {room.player_names[i]} is lost [{getframeinfo(currentframe()).lineno}]")
        if room.room_id in self.game_dict:
            self.room_cnt -= 1
            self.game_dict.pop(room.room_id)
        self.my_logger.warning(f"Connection issue to player(s): "
                               f"{room.player_names}, game room is being closed!")

    def remove_player1(self, room):
        self.cnt -= 1
        room.player_names.pop()  # player1
        room.player_readers.pop()  # player1
        room.player_writers.pop()  # player1
        room.choosable = True

    async def check_read_room(self, room: RoomState):
        # used for read() after the connected player is in a game room
        string_lists = [[], []]
        connected = CONNECTED

        for i in range(2):
            quit_flg = "not set"  # if "q", player quit the game normally, don't use "1", used for selection ready
            try:
                # receiving event strings or game setting strings
                received = await room.player_readers[i].readuntil(separator=b"AB")
                info = list(received.decode()[:-2])
                quit_flg = info[0]
                """
                In shooter-chopper game, player_role_ids is [0, 1] or [1, 0] string_lists[0] is always shooter,
                room.player_role_ids[i] actually represents the events for the role
                In shooter-shooter or chopper-chopper game, since player_role_ids is [0, 0] or [1, 1]
                room.player_role_ids[i] can't be directly used as the index for string_lists[]
                use i directly to assign the received info to string_lists[]
                """
                # (1 - room.player_role_ids[i]) makes it right as the index for string_lists[]
                # if len(string_lists[room.player_role_ids[i]]) > 0:
                #     # 1-1 = 0, 1-0 = 1
                #     index = 1 - room.player_role_ids[i]
                #     string_lists[index] = list(received.decode()[:-2])  # discard trailing "AB"
                if room.player_role_ids[0] == room.player_role_ids[1]:
                    string_lists[i] = info
                else:
                    string_lists[room.player_role_ids[i]] = info
            except (ConnectionError, asyncio.IncompleteReadError):
                # writer.close() is not needed because the writer is already closed
                # await writer.wait_closed(): this line could never be returned when the writer is already closed
                self.my_logger.warning(
                    f"Connection to player <{room.player_names[i]}> is lost [{getframeinfo(currentframe()).lineno}]")
                connected = not CONNECTED
            else:  # if the connection is properly connected, there will be no ConnectionError exception raised
                if quit_flg == QUIT:  # player quit the game normally, don't use "1", used for selection ready
                    self.my_logger.warning(
                        f"Player <{room.player_names[i]}> quit the game")
                    if not room.player_writers[i].is_closing():
                        room.player_writers[i].close()
                        await room.player_writers[i].wait_closed()
                    # self.remove_room(room, i)
                    connected = not CONNECTED
            if not connected:  # if player i is disconnected
                i = (i + 1) % 2  # if i = 0, make it 1, if i = 1, make it 0
                room.player_writers[i].write("DisconnectedAB".encode())
                room.player_writers[i].close()
                await room.player_writers[i].wait_closed()
                self.my_logger.warning(
                    f"Player <{room.player_names[i]}> is disconnected [{getframeinfo(currentframe()).lineno}]")
                return connected, string_lists
        return connected, string_lists

    def create_room(self, player_name, reader, writer):
        self.room_cnt += 1
        self.room_id += 1
        room = RoomState()
        room.room_id = self.room_id
        room.player_names.append(player_name)  # player0
        room.player_readers.append(reader)  # player0
        room.player_writers.append(writer)  # player0
        # room.player_role_ids.append(0)  # player0
        self.game_dict[self.room_id] = room
        return room

    async def join(self, player_name, reader, writer):
        # if the client request to join an existing game, a new coroutine will be created with this function
        chosen_room_id = None
        while True:
            # sending room list to the client, no matter full or not [[player0_name, game_ready, room_id],]
            # if game_dict is empty, an empty list will be assigned to "rooms_lst" and no exception will be raised
            rooms_lst = [[self.game_dict[room_id].player_names[0], self.game_dict[room_id].check_ready(), room_id]
                         for room_id in [*self.game_dict] if self.game_dict[room_id].choosable]
            if not rooms_lst:
                rooms_lst = [["no game", False, 0]]
            rooms_lst_enc = (json.dumps(rooms_lst) + "AB").encode()
            writer.write(rooms_lst_enc)
            r = await self.check_read(player_name, reader, writer, False)  # return connected, string
            if not r[0]:
                return not CONNECTED, chosen_room_id
            else:
                info = list(r[1].split(","))
                if int(info[0]) in [*self.game_dict]:  # r1:string = "room_id, role_id"
                    writer.write("okAB".encode())
                    chosen_room_id = int(info[0])
                    chosen_role_id = int(info[1])
                    # role_id = int(r[1][2])
                    r = await self.check_read(player_name, reader, writer, False)  # r/w cycle
                    if not r[0]:
                        return not CONNECTED, chosen_room_id
                    """
                    # self.game_dict[chosen_room_id].player_joined = True
                    player_joined can't be set true here, otherwise player0 will start to write to player1's writer too
                    soon to cause the "mark 1" read() below read more bytes in the buffer from player1' reader because the
                    player1 client writes the "events" info. Best way could be using reader.readuntil() with a seperator
                    letter (e.g. "AB"), tested working without turn on/off confirmation.
                    """
                    self.game_dict[chosen_room_id].choosable = False
                    self.game_dict[chosen_room_id].player_names.append(player_name)  # player1
                    self.game_dict[chosen_room_id].player_readers.append(reader)  # player1
                    self.game_dict[chosen_room_id].player_writers.append(writer)  # player1
                    self.game_dict[chosen_room_id].player_role_ids[1] = chosen_role_id  # player1
                    chosen_room = self.game_dict[chosen_room_id]
                    while not chosen_room.game_set0:  # meaning player0 has not finished setting game yet
                        # clock.tick(FPS)
                        writer.write(
                            f"{self.game_dict[chosen_room_id].player_names[0]} is not ready yet, please wait...AB".encode())
                        r = await self.check_read(player_name, reader, writer, True)  # return connected, string
                        if not r[0]:
                            if chosen_room_id in self.game_dict:
                                self.remove_player1(self.game_dict[chosen_room_id])
                            return not CONNECTED, chosen_room_id
                    writer.write("Game is ready to playAB".encode())
                    r = await self.check_read(player_name, reader, writer, True)  # (mark 1) return connected, string
                    if not r[0]:
                        if chosen_room_id in self.game_dict:
                            self.remove_player1(self.game_dict[chosen_room_id])
                        return not CONNECTED, chosen_room_id
                    return CONNECTED, chosen_room

    def new_connection(self, player_name):
        self.cnt += 1
        self.client_id += 1
        self.my_logger.info(f"New connection to player '{player_name}', client id = {self.client_id}, "
                            f"total connections: {self.cnt}")

    async def new_client(self, reader, writer):
        # any new client connection goes here first
        client_task = asyncio.current_task()
        client_name = ""
        player_id = 0  # player_id will be assigned according to conn_type: create - 0, join - 1
        player_info = None
        room = None
        player_name = None

        try:  # initial read() needs different handling than self.check_read
            received = await reader.readuntil(separator=b"AB")
            temp = received.decode()[:-2] # "handshake,Amy" = ["handshake", "Amy"]
            player_info = temp.split(",")  # "{conn_type},{self.player_name}"
        except (ConnectionError, asyncio.IncompleteReadError):
            self.my_logger.warning(
                f"Connection to player <{player_name}> is lost [{getframeinfo(currentframe()).lineno}]")
        # if the connection is properly closed, there will be no ConnectionError exception raised
        # but empty byte will be received
        else:
            if not received:
                self.my_logger.warning(
                    f"Connection to player <{player_name}> is lost [{getframeinfo(currentframe()).lineno}]")
                if not writer.is_closing():
                    writer.close()
                    await writer.wait_closed()

        if player_info[0] != "handshake":  # if not "handshake", it's an invalid connection request
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()
            return
        elif player_info[0] == "handshake":  # handshake connection, waiting for conn_type
            player_name = player_info[1]
            writer.write("okAB".encode())  # for client to confirm successful handshake

            try:  # initial read() needs different handling than self.check_read, this is just for r/w cycle
                received = await reader.readuntil(separator=b"AB")
            except (ConnectionError, asyncio.IncompleteReadError):
                self.my_logger.warning(
                    f"Connection to player <{player_name}> is lost [{getframeinfo(currentframe()).lineno}]")
            else:  # if the connection is properly closed, there will be no ConnectionError exception raised
                if not received:
                    self.my_logger.warning(
                        f"Connection to player <{player_name}> is lost [{getframeinfo(currentframe()).lineno}]")
                    if not writer.is_closing():
                        writer.close()
                        await writer.wait_closed()

            self.new_connection(player_name)
            client_name = f"{player_name}-{self.client_id}"
            client_task.set_name(client_name)
            writer.write((str(self.client_id) + "AB").encode())

            r = await self.check_read(player_name, reader, writer, False)  # waiting for "create" or "join"
            if not r[0]:  # return connected, string
                return
            else:
                player_info = r[1].split(",")

            conn_type = player_info[0]  # client reply: f"{conn_type},{player_name}"
            if conn_type == "create":
                player_id = 0
                room = self.create_room(player_name, reader, writer)
                room.player_task_names.append(client_name)  # player0
            elif conn_type == "join":
                player_id = 1
                try:
                    join_room = await self.join(player_info[1], reader, writer)  # return CONNECTED, choice
                except ConnectionError:
                    # self.my_logger.warning(f"Connection to player {player_name} is lost")
                    return

                if not join_room[0]:
                    self.remove_player1(join_room[1])
                    return
                else:
                    room = join_room[1]
                    room.player_task_names.append(client_name)  # player1
                    # room.player_role_ids[1] = 1  # player1
                    room.player_joined = True
                    """
                    player_joined can't be set to True in self.join(), otherwise player0 will enter the
                    routine game loop too soon, refer to comments in self.join() for details
                    """
        # start = pygame.time.get_ticks()
        # clock = pygame.time.Clock()
        while True:  # wait for the 2nd player to join the room and the 1st player to set the game mode
            # clock.tick()  # needs to be called for get_fps() to result in correct value
            # tick = pygame.time.get_ticks() - start
            """
            once there are two players in one game room, game_ready for that room is set to
            True, then the task for player_1 who joins the room after player_0 created the room
            will be returned so that the server game tick info to both player_0 and player_1
            can be handled by the task created for player_0
            """
            if room.check_ready():
                if player_id == 1:  # player_id = 1 means this is the task for player_1
                    self.my_logger.warning(
                        f"Player <{room.player_names[1]}> joined player <{room.player_names[0]}>'s game")
                    return  # the task (for player_1) is returned (completed) once player_1 is in the room
                # data = "Game Ready".encode()
                room.running = True
                room.player_writers[0].write(
                    f"Game Ready,{room.player_names[1]},{room.map_id},{room.match_id},"
                    f"{room.player_role_ids[0]},{room.player_role_ids[1]}AB".encode())
                room.player_writers[1].write(
                    f"Game Ready,{room.player_names[0]},{room.map_id},{room.match_id},"
                    f"{room.player_role_ids[0]},{room.player_role_ids[1]}AB".encode())
                break  # break current while loop to start the routine game tick
            else:  # if game_ready is False, it means there is only player_0 in the game room
                data = f"New game is created, waiting for the second player to join...AB".encode()
                room.player_writers[0].write(data)
                r = await self.check_read(
                    room.player_names[0], room.player_readers[0],
                    room.player_writers[0], True)  # r/w cycle and check the connection
                if not r[0]:  # return connected, string
                    if room.room_id in self.game_dict and player_id == 0:
                        self.cnt -= 1
                        self.room_cnt -= 1
                        self.game_dict.pop(room.room_id)
                    else:
                        self.remove_player1(room)
                        room.player_task_names.pop()  # player1
                        room.player_role_ids.pop()  # player1
                        room.player_joined = False
                    return
                else:
                    info = list(r[1])
                    if info[0] == "1":  # ready, map_id, match_id, role_id
                        room.map_id = info[1]
                        room.match_id = info[2]
                        room.player_role_ids[0] = int(info[3])  # role_id for player0
                        # room.player_role_ids[1] = 0
                        # room.player_role_ids[1] = 1 - room.player_role_ids[0]  # 1-1 = 0, 1-0 = 1
                        room.game_set0 = True
            # print((perf_counter() - start)*1000)

        # game_type = game_class_s.game_type_lst[room.player_role_ids[0] + room.player_role_ids[1]]
        # g = game_type(
        #     self.screen, SCREEN_WIDTH, SCREEN_HEIGHT, int(room.map_id), 0,
        #     int(room.match_id))  # map_id, level_id, match_id

        while room.running:
            g = self.set_game(room)
            await self.game(room, g)  # this is the routine game tick
            if room.running:
                await self.re_select(room)
        self.remove_room(room)

    def set_game(self, room):
        game_type = game_class_s.game_type_lst[room.player_role_ids[0] + room.player_role_ids[1]]
        g = game_type(
            self.screen, SCREEN_WIDTH, SCREEN_HEIGHT, int(room.map_id), 0,
            int(room.match_id))  # map_id, level_id, match_id
        return g

    async def re_select(self, room: RoomState):
        room.winner = "nobody"
        room.game_set0 = False
        room.player_joined = False
        start = pygame.time.get_ticks()
        clock = pygame.time.Clock()
        while not room.check_ready():
            clock.tick()  # needs to be called for get_fps() to result in correct value
            tick = pygame.time.get_ticks() - start
            if tick < FPS_T:
                # print(FPS_T - tick)
                await asyncio.sleep((FPS_T - tick) / 1000)
            start = pygame.time.get_ticks()

            r = await self.check_read(room.player_names[0], room.player_readers[0], room.player_writers[0], True)
            if not r[0]:
                room.running = False
                return
            else:
                if r[1][0] == "1":  # ready, map_id, match_id, role_id
                    room.map_id = r[1][1]
                    room.match_id = r[1][2]
                    room.player_role_ids[0] = int(r[1][3])  # role_id for player0
                    # room.player_role_ids[1] = 1 - room.player_role_ids[0]  # 1-1 = 0, 1-0 = 1
                    room.game_set0 = True

            r = await self.check_read(room.player_names[1], room.player_readers[1], room.player_writers[1], True)
            if not r[0]:
                room.running = False
                return
            else:
                if r[1][0] == "1":  # ready, role_id, 0, 0
                    room.player_role_ids[1] = int(r[1][1])  # role_id for player0
                    room.player_joined = True
                    # player1_resel_status = "okAB"

            room.player_writers[0].write(f"re-selectingAB".encode())
            room.player_writers[1].write(f"re-selectingAB".encode())

        r = await self.check_read_room(room)  # r/w cycle
        if not r[0]:
            room.running = False
            return

        room.player_writers[0].write(
            f"Game Ready,{room.player_names[1]},{room.map_id},{room.match_id},"
            f"{room.player_role_ids[0]},{room.player_role_ids[1]}AB".encode())
        room.player_writers[1].write(
            f"Game Ready,{room.player_names[0]},{room.map_id},{room.match_id},"
            f"{room.player_role_ids[0]},{room.player_role_ids[1]}AB".encode())

    async def game(self, room, g):
        # g.new()
        # g.map_id = int(room.map_id)
        g.current_level_no = int(room.map_id)
        g.match_id = int(room.match_id)
        g.new()
        # gs = GameState()
        self.my_logger.warning(f"<{room.player_names[0]}> is gaming with <{room.player_names[1]}>!")
        # start = pygame.time.get_ticks()
        # start_p = perf_counter()
        clock = pygame.time.Clock()
        while g.playing:  # This is the routine game tick
            clock.tick()  # needs to be called for get_fps() to result in correct value
            # tick = pygame.time.get_ticks() - start
            # print(FPS_T - tick)
            # if tick < FPS_T:
            #     await asyncio.sleep((FPS_T - tick)/1000)
            # start = pygame.time.get_ticks()
            actual_fps = int(clock.get_fps())
            # print(actual_fps)
            actual_tick = int(pygame.time.get_ticks())
            if 0 < actual_fps <= 40 and (actual_tick - self.timer) > 5000:
                self.timer = actual_tick
                self.my_logger.warning(f"Low fps: {actual_fps}")
            # print(room.room_id, (perf_counter()-start_p)*1000)
            # print(perf_counter()-start_p)
            r = await self.check_read_room(room)
            # for i in range(2):
            #     for j in range(7):
            #         if r[1][i][j] == "1":
            #             print(r[1])
            #             break
            # start_p = perf_counter()
            if not r[0]:
                room.running = False
                return
            g.events_lst0 = r[1][0]
            g.events_lst1 = r[1][1]

            if g.events_lst0[0] != HOLD and g.events_lst1[0] != HOLD:
                if not g.events():
                    room.running = False
                    return
                g.update()

            byte = (json.dumps(g.gs_conversion()).encode())
            compressed_send_bytes = compress(byte) + b"AB"
            # print(compressed_send_bytes)
            room.player_writers[0].write(compressed_send_bytes)
            room.player_writers[1].write(compressed_send_bytes)

            if g.new_round:
                g.restart()
                g.new_round = False

        room.winner = g.winner

    def handle_exception(self, loop, context):
        # context["message"] will always be there; but context["exception"] may not
        # msg = context.get("exception", context["message"])
        self.my_logger.error(f"Caught exception: {context}")

    async def main(self, host, port):
        server = await asyncio.start_server(self.new_client, host, port, )
        self.loop = server.get_loop()
        self.loop.set_exception_handler(self.handle_exception)
        self.my_logger.info(f"Server started at {host}:{port}")
        await server.serve_forever()

if __name__ == "__main__":
    s = Server()
    s.my_logger.info("----------------------New log started---------------------------------")
    try:
        # No IP address or "0.0.0.0" is provided to allow the server to listen to all connections on the port
        asyncio.run(s.main('0.0.0.0', 8887))
    except KeyboardInterrupt:
        s.my_logger.warning("Server terminated by player.")


"""
This module takes a single input file and returns the final x, y positions with the coins collected
First line has the board dimensions, the second line has the intial starting postion
Third line has the movements in N, E, W, S
Fourth line has the restricted x, y positions

The module returns [-1,-1,0] if there are any file validation failures

"""

__author__ = "Saikumar Beera"

def fCalculcation (vDirections, vMoves, x, y, rx, ry) :
    """This function does the edge validiation, calculation of the
    valid moves and the total coins collected.
    The function returns the x , y and total coins collections
    """
    if x > rx or y > ry:
        return (-1, -1, 0)
    vValidStack = []
    for  d in vDirections:
        x, y = x+d[0], y+d[1]
        if [x,y] not in vValidStack:
            vValidStack.append([x,y])

        if [x,y] in vMoves:
            x, y = x-d[0], y-d[1]
            vValidStack.pop()
            continue
        if  x > rx or x < 0 or  y > ry or y < 0 :
            x, y = x-d[0], y-d[1]
            vValidStack.pop()
            continue
    return [x, y, len(vValidStack)]

def pacman(inputFile):
    """ Use this function to format your input/output arguments. Be sure not to change the order of the output arguments.
    Remember that code organization is very important to us, so we encourage the use of helper fuctions and classes as you see fit.

    Input:
        1. input_file (String) = contains the name of a text file you need to read that is in the same directory, includes the ".txt" extension
           (ie. "input.txt")
    Outputs:
        1. final_pos_x (int) = final x location of Pacman
        2. final_pos_y (int) = final y location of Pacman
        3. coins_collected (int) = the number of coins that have been collected by Pacman across all movements
    """

    with open (inputFile) as f:
        plines = f.read().splitlines()

    pCoins = 0
    # Directions stored
    lDirections = [[-1, 0], [0, -1], [0, 1], [1, 0]]

    board_x, board_y = list(map(int, plines[0].split(' ')))
    initial_x, initial_y = list(map(int, plines[1].split(' ')))
    iDirections = list(plines[2])
    iMoves = list(plines[3:])

    # Our transfromations

    vMoves = []
    vDirections = []

    for i in range(len(iMoves)):
        vMoves.append(list(map(int, iMoves[i].split(' '))))

    for j in range(len(iDirections)):
        if iDirections[j]=='N':
            vDirections.append(lDirections[2])
        if iDirections[j]=='S':
            vDirections.append(lDirections[1])
        if iDirections[j]=='E':
            vDirections.append(lDirections[3])
        if iDirections[j]=='W':
            vDirections.append(lDirections[0])

    return fCalculcation(vDirections, vMoves, initial_x, initial_y, board_x, board_y)


import os
import firebase_admin
from firebase_admin import credentials
from firebase_admin import firestore
from dotenv import load_dotenv
from interfaces.db_interface import IDatabase
from datetime import datetime

load_dotenv()

config = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
project_id = os.environ.get('project_id')

# Apply credentials
# cred = credentials.Certificate(config)

# firebase_admin.initialize_app(cred)

cred = credentials.ApplicationDefault()

firebase_admin.initialize_app(
    cred, {"projectId": project_id}
)

db = firestore.client()

class FireStoreDB(IDatabase):

    __instance = None

    def __init__(self):
        FireStoreDB.connect(self)

    @staticmethod
    def get_instance(re_init=False):
        if FireStoreDB.__instance is None or re_init is True:
            return FireStoreDB()
        return FireStoreDB.__instance

    def connect(self):
        print("Connecting to Firestore")
        self.odds_ref = db.collection('odds')
        return self.odds_ref

    # def disconnect(self):
    #     print("Disconnecting from Firestore database")
    #     return True

    def create(
        self,
        league: str,
        home_team: str,
        away_team: str,
        home_team_win_odds: float,
        away_team_win_odds: float,
        draw_odds: float,
        game_date
    ):
        print("-Saving odds in Firestore")

        try:
            odd = {
                'league': league,
                'home_team': home_team,
                'away_team': away_team,
                'home_team_win_odds': home_team_win_odds,
                'away_team_win_odds': away_team_win_odds,
                'draw_odds': draw_odds,
                'game_date': datetime.strptime(game_date.strip(), "%Y-%m-%d")
            }

            saved_odds = self.odds_ref.add(odd)
            print(saved_odds[1].id)
            added_odds = self.odds_ref.document(
                    saved_odds[1].id
                ).get().to_dict()
            # add unique identifier of document as id
            added_odds['id'] = saved_odds[1].id
            reason = "Successfully created"
            print(reason)
            return True, added_odds
        except Exception as e:
            reason = (
                "-Failed to create odds in Firestore: "
                + f"{type(e).__name__} {str(e)}"
            )
            print(reason)
            return False, reason

    def read(self, league, date_from, date_to):
        print("-Viewing odds in database")
        # read odds from database
        date_from = datetime.strptime(date_from.strip(), "%Y-%m-%d")
        date_to = datetime.strptime(date_to.strip(), "%Y-%m-%d")
        read_odds = []
        try:
            docs = self.odds_ref.where(
                'league', '==', league).where(
                'game_date', '>=', date_from).where(
                'game_date', '<=', date_to).stream()

            for doc in docs:
                print('{} => {} '.format(doc.id, doc.to_dict()))
                read_odds.append(doc.to_dict())
            reason = "-Odds viewed successfully"
            print(reason)
            return True, read_odds
        except Exception as e:
            reason = (
                f"-Failed to read odds - {e}"
            )
            print(reason)
            return False, reason, ""

    def update(
        self,
        odd_id,
        league,
        home_team,
        away_team,
        home_team_win_odds,
        away_team_win_odds,
        draw_odds,
        game_date
    ):
        print("-Updating odds ")
        # update odds from database
        try:
            self.odds_ref.document(odd_id).set({
                "league": league,
                "home_team": home_team,
                "away_team": away_team,
                "home_team_win_odds": home_team_win_odds,
                "away_team_win_odds": away_team_win_odds,
                "draw_odds": draw_odds,
                "game_date": datetime.strptime(game_date.strip(), "%Y-%m-%d")
            })

            reason = "-Odds updated successfully"
            print(reason)
            returned, odd = FireStoreDB.get_odd(self, odd_id)
            return True, odd
        except Exception as e:
            reason = (
                f"-Failed to update odds, error- {e} "
            )
            print(reason)
            return False, reason

    def delete(self, league, home_team, away_team, game_date):
        print("-Deleting odds from database")

        try:
            game_date = datetime.strptime(game_date.strip(), "%Y-%m-%d")
            docs = self.odds_ref.where(
                'league', '==', league).where(
                'home_team', '==', home_team).where(
                'away_team', '==', away_team).where(
                'game_date', '==', game_date).get()
            print(docs)
            if docs:
                for doc in docs:
                    self.odds_ref.document(doc.id).delete()

                reason = "-Odds deleted successfully"
                return True, reason

            return False, "Odds not found"
        except Exception as e:
            reason = (
                f"-Failed to delete odds - {e}"
            )
            print(reason)
            return False, reason

    def get_odd(self, odd_id):
        try:
            doc = self.odds_ref.document(odd_id).get().to_dict()
            if doc:
                print(doc)
                return True, [doc]
            return False, "Odds not found"
        except Exception as e:
            return False, f"An Error Occured: {e}"

    def find_by_field(self, league, home_team, away_team, game_date):
        """find specific field"""
        game_date = datetime.strptime(game_date.strip(), "%Y-%m-%d")
        try:
            returned_odds = self.odds_ref.where(
                    'league', '==', league).where(
                    'home_team', '==', home_team).where(
                    'away_team', '==', away_team).where(
                    'game_date', '==', game_date).get()
            print(returned_odds)
            if returned_odds:
                return True, returned_odds
            return False, "couldn't find odds"
        except Exception as e:
            return False, f"failed to read Odds from db -{e}"


#adapted with appreciation from https://github.com/jump-cellpainting/pilot-cpjump1-analysis

from multiprocessing.sharedctypes import Value
import os
import random
import textwrap

import seaborn as sns
import pandas as pd
import numpy as np
import scipy
from sklearn.utils.validation import check_is_fitted
from sklearn.utils import check_array, as_float_array
from sklearn.base import TransformerMixin, BaseEstimator
import kneed
import matplotlib.pyplot as plt
import seaborn
import ast
import skimage
import boto3
from botocore.handlers import disable_signing

random.seed(9000)

def get_metacols(df):
    """return a list of metadata columns"""
    return [c for c in df.columns if c.startswith("Metadata_")]

def get_featurecols(df):
    """returna  list of featuredata columns"""
    return [c for c in df.columns if not c.startswith("Metadata")]

def get_metadata(df):
    """return dataframe of just metadata columns"""
    return df[get_metacols(df)]

def get_featuredata(df):
    """return dataframe of just featuredata columns"""
    return df[get_featurecols(df)]

def remove_negcon_empty_wells(df):
    """return dataframe of non-negative-control wells"""
    df = (
        df.query('Metadata_control_type!="negcon"')
        .dropna(subset=['Metadata_broad_sample'])
        .reset_index(drop=True)
    )
    return df

def percent_score(null_dist, corr_dist, how='right'):
    """
    Calculates the Percent replicating
    :param null_dist: Null distribution
    :param corr_dist: Correlation distribution
    :param how: "left", "right" or "both" for using the 5th percentile, 95th percentile or both thresholds
    :return: proportion of correlation distribution beyond the threshold
    """
    if how == 'right':
        perc_95 = np.nanpercentile(null_dist, 95)
        above_threshold = corr_dist > perc_95
        return 100 * np.mean(above_threshold.astype(float)), perc_95
    if how == 'left':
        perc_5 = np.nanpercentile(null_dist, 5)
        below_threshold = corr_dist < perc_5
        return 100 * np.mean(below_threshold.astype(float)), perc_5
    if how == 'both':
        perc_95 = np.nanpercentile(null_dist, 95)
        above_threshold = corr_dist > perc_95
        perc_5 = np.nanpercentile(null_dist, 5)
        below_threshold = corr_dist < perc_5
        return 100 * (np.mean(above_threshold.astype(float)) + np.mean(below_threshold.astype(float))), perc_95, perc_5

def corr_between_replicates(df, group_by_feature):
    """
    Correlation between replicates
    :param df: pd.DataFrame
    :param group_by_feature: Feature name to group the data frame by
    :return: list-like of correlation values
    """
    replicate_corr = []
    replicate_grouped = df.groupby(group_by_feature)
    for name, group in replicate_grouped:
        group_features = get_featuredata(group)
        corr = np.corrcoef(group_features)
        if len(group_features) == 1:  # If there is only one replicate on a plate
            replicate_corr.append(np.nan)
        else:
            np.fill_diagonal(corr, np.nan)
            replicate_corr.append(np.nanmedian(corr))  # median replicate correlation
    return replicate_corr

def corr_between_non_replicates(df, n_samples, n_replicates, metadata_compound_name):
    """
    Null distribution between random "replicates".
    :param df: pandas.DataFrame
    :param n_samples: int
    :param n_replicates: int
    :param metadata_compound_name: Compound name feature
    :return: list-like of correlation values, with a  length of `n_samples`
    """
    df.reset_index(drop=True, inplace=True)
    null_corr = []
    while len(null_corr) < n_samples:
        compounds = random.choices([_ for _ in range(len(df))], k=n_replicates)
        sample = df.loc[compounds].copy()
        if len(sample[metadata_compound_name].unique()) == n_replicates:
            sample_features = get_featuredata(sample)
            corr = np.corrcoef(sample_features)
            np.fill_diagonal(corr, np.nan)
            null_corr.append(np.nanmedian(corr))  # median replicate correlation
    return null_corr

def corr_between_replicates_across_plates(df, reference_df, pertcol = 'Metadata_pert_iname'):
    items = list(df[pertcol].unique())
    common_columns = [x for x in df.columns if x in reference_df.columns]
    df = df[common_columns]
    reference_df = reference_df[common_columns]
    replicate_corr = []
    for item in items:
        compound_df = df.query(pertcol+' == @item')
        compound_reference_df = reference_df.query(pertcol+' == @item')

        compound_df_profiles = get_featuredata(compound_df).values
        compound_reference_df_profiles = get_featuredata(compound_reference_df).values
        try:
            corr = np.corrcoef(compound_df_profiles, compound_reference_df_profiles)
            corr = corr[0:len(compound_df_profiles), len(compound_df_profiles):]

            corr_median_value = np.nanmedian(corr, axis=1)
            corr_median_value = np.nanmedian(corr_median_value)

            replicate_corr.append(corr_median_value)
        except:
            print(item,corr)
    return replicate_corr

def corr_between_non_replicates_across_plates(df, reference_df, n_samples, pertcol = 'Metadata_pert_iname'):
    np.random.seed(9000)
    common_columns = [x for x in df.columns if x in reference_df.columns]
    df = df[common_columns]
    reference_df = reference_df[common_columns]
    null_corr = []
    compounds = list(df[pertcol].unique())
    while len(null_corr) < n_samples:
        both_compounds = np.random.choice(compounds, size=2, replace=False)
        compound1 = both_compounds[0]
        compound2 = both_compounds[1]

        compound1_df = df.query(pertcol+' == @compound1')
        compound2_df = reference_df.query(pertcol+' == @compound2')

        compound1_df_profiles = get_featuredata(compound1_df).values
        compound2_df_profiles = get_featuredata(compound2_df).values
        try:
            corr = np.corrcoef(compound1_df_profiles, compound2_df_profiles)
            corr = corr[0:len(compound1_df_profiles), len(compound1_df_profiles):]

            corr_median_value = np.nanmedian(corr, axis=1)
            corr_median_value = np.nanmedian(corr_median_value)

            null_corr.append(corr_median_value)
        except:
            pass

    return null_corr

def corr_between_compound_moa(df, metadata_moa, metadata_compound_name):
    """
        Correlation between compounds with the same MOA
        Parameters:
        -----------
        df: pd.DataFrame
        metadata_moa: MOA feature
        metadata_compound_name: Compound name feature
        Returns:
        --------
        list-like of correlation values
     """
    replicate_corr = []

    profile_df = (
        get_metadata(df)
        .assign(profiles=list(get_featuredata(df).values))
    )

    replicate_grouped = (
        profile_df.groupby([metadata_moa, metadata_compound_name]).profiles
            .apply(list)
            .reset_index()
    )

    moa_grouped = (
        replicate_grouped.groupby([metadata_moa]).profiles
            .apply(list)
            .reset_index()
    )

    for i in range(len(moa_grouped)):
        if len(moa_grouped.iloc[i].profiles) > 1:
            compound1_profiles = moa_grouped.iloc[i].profiles[0]
            compound2_profiles = moa_grouped.iloc[i].profiles[1]

            corr = np.corrcoef(compound1_profiles, compound2_profiles)
            corr = corr[0:len(moa_grouped.iloc[i].profiles[0]), len(moa_grouped.iloc[i].profiles[0]):]
            # np.fill_diagonal(corr, np.nan)
            replicate_corr.append(np.nanmedian(corr))

    return replicate_corr

def null_corr_between_compound_moa(df, n_samples, metadata_moa, metadata_compound_name):
    """
        Null distribution between random pairs of compounds.
        Parameters:
        ------------
        df: pandas.DataFrame
        n_samples: int
        metadata_moa: MOA feature
        metadata_compound_name: Compound name feature
        Returns:
        --------
        list-like of correlation values, with a  length of `n_samples`
    """
    df.reset_index(drop=True, inplace=True)
    null_corr = []

    profile_df = (
        get_metadata(df)
        .assign(profiles=list(get_featuredata(df).values))
    )

    replicate_grouped = (
        profile_df.groupby([metadata_moa, metadata_compound_name]).profiles
            .apply(list)
            .reset_index()
    )

    while len(null_corr) < n_samples:
        compounds = random.choices([_ for _ in range(len(replicate_grouped))], k=2)
        compound1_moa = replicate_grouped.iloc[compounds[0]].Metadata_moa
        compound2_moa = replicate_grouped.iloc[compounds[1]].Metadata_moa
        if compound1_moa != compound2_moa:
            compound1_profiles = replicate_grouped.iloc[compounds[0]].profiles
            compound2_profiles = replicate_grouped.iloc[compounds[1]].profiles
            corr = np.corrcoef(compound1_profiles, compound2_profiles)
            corr = corr[0:len(replicate_grouped.iloc[0].profiles), len(replicate_grouped.iloc[0].profiles):]
            # np.fill_diagonal(corr, np.nan)
            null_corr.append(np.nanmedian(corr))  # median replicate correlation
    return null_corr

def correlation_between_modalities(modality_1_df, modality_2_df, modality_1, modality_2, metadata_common, metadata_perturbation):
    """
    Compute the correlation between two different modalities.
    :param modality_1_df: Profiles of the first modality
    :param modality_2_df: Profiles of the second modality
    :param modality_1: feature that identifies perturbation pairs
    :param modality_2: perturbation name feature
    :param metadata_common: perturbation name feature
    :param metadata_perturbation: perturbation name feature
    :return: list-like of correlation values
    """
    list_common_perturbation_groups = list(np.intersect1d(list(modality_1_df[metadata_common]), list(modality_2_df[metadata_common])))

    merged_df = pd.concat([modality_1_df, modality_2_df], ignore_index=False, join='inner')

    modality_1_df = merged_df.query('Metadata_modality==@modality_1')
    modality_2_df = merged_df.query('Metadata_modality==@modality_2')

    corr_modalities = []

    for group in list_common_perturbation_groups:
        modality_1_perturbation_df = modality_1_df.loc[modality_1_df[metadata_common] == group]
        modality_2_perturbation_df = modality_2_df.loc[modality_2_df[metadata_common] == group]

        for sample_1 in modality_1_perturbation_df[metadata_perturbation].unique():
            for sample_2 in modality_2_perturbation_df[metadata_perturbation].unique():
                modality_1_perturbation_sample_df = modality_1_perturbation_df.loc[modality_1_perturbation_df[metadata_perturbation] == sample_1]
                modality_2_perturbation_sample_df = modality_2_perturbation_df.loc[modality_2_perturbation_df[metadata_perturbation] == sample_2]

                modality_1_perturbation_profiles = get_featuredata(modality_1_perturbation_sample_df)
                modality_2_perturbation_profiles = get_featuredata(modality_2_perturbation_sample_df)

                corr = np.corrcoef(modality_1_perturbation_profiles, modality_2_perturbation_profiles)
                corr = corr[0:len(modality_1_perturbation_profiles), len(modality_1_perturbation_profiles):]
                corr_modalities.append(np.nanmedian(corr))  # median replicate correlation

    return corr_modalities

def null_correlation_between_modalities(modality_1_df, modality_2_df, modality_1, modality_2, metadata_common, metadata_perturbation, n_samples):
    """
    Compute the correlation between two different modalities.
    :param modality_1_df: Profiles of the first modality
    :param modality_2_df: Profiles of the second modality
    :param modality_1: "Compound", "ORF" or "CRISPR"
    :param modality_2: "Compound", "ORF" or "CRISPR"
    :param metadata_common: feature that identifies perturbation pairs
    :param metadata_perturbation: perturbation name feature
    :param n_samples: int
    :return:
    """
    list_common_perturbation_groups = list(np.intersect1d(list(modality_1_df[metadata_common]), list(modality_2_df[metadata_common])))

    merged_df = pd.concat([modality_1_df, modality_2_df], ignore_index=False, join='inner')

    modality_1_df = merged_df.query('Metadata_modality==@modality_1')
    modality_2_df = merged_df.query('Metadata_modality==@modality_2')

    null_modalities = []

    count = 0

    while count < n_samples:
        perturbations = random.choices(list_common_perturbation_groups, k=2)
        modality_1_perturbation_df = modality_1_df.loc[modality_1_df[metadata_common] == perturbations[0]]
        modality_2_perturbation_df = modality_2_df.loc[modality_2_df[metadata_common] == perturbations[1]]

        for sample_1 in modality_1_perturbation_df[metadata_perturbation].unique():
            for sample_2 in modality_2_perturbation_df[metadata_perturbation].unique():
                modality_1_perturbation_sample_df = modality_1_perturbation_df.loc[modality_1_perturbation_df[metadata_perturbation] == sample_1]
                modality_2_perturbation_sample_df = modality_2_perturbation_df.loc[modality_2_perturbation_df[metadata_perturbation] == sample_2]

                modality_1_perturbation_profiles = get_featuredata(modality_1_perturbation_sample_df)
                modality_2_perturbation_profiles = get_featuredata(modality_2_perturbation_sample_df)

                corr = np.corrcoef(modality_1_perturbation_profiles, modality_2_perturbation_profiles)
                corr = corr[0:len(modality_1_perturbation_profiles), len(modality_1_perturbation_profiles):]
                null_modalities.append(np.nanmedian(corr))  # median replicate correlation
        count += 1

    return null_modalities

class ZCA_corr(BaseEstimator, TransformerMixin):
    def __init__(self, copy=False):
        self.copy = copy

    def estimate_regularization(self, eigenvalue):
        x = [_ for _ in range(len(eigenvalue))]
        kneedle = kneed.KneeLocator(x, eigenvalue, S=1.0, curve='convex', direction='decreasing')
        reg = eigenvalue[kneedle.elbow]/10.0
        return reg # The complex part of the eigenvalue is ignored

    def fit(self, X, y=None):
        """
        Compute the mean, sphereing and desphereing matrices.
        Parameters
        ----------
        X : array-like with shape [n_samples, n_features]
            The data used to compute the mean, sphereing and desphereing
            matrices.
        """
        X = check_array(X, accept_sparse=False, copy=self.copy, ensure_2d=True)
        X = as_float_array(X, copy=self.copy)
        self.mean_ = X.mean(axis=0)
        X_ = X - self.mean_
        cov = np.dot(X_.T, X_) / (X_.shape[0] - 1)
        V = np.diag(cov)
        df = pd.DataFrame(X_)
        corr = np.nan_to_num(df.corr()) # replacing nan with 0 and inf with large values
        G, T, _ = scipy.linalg.svd(corr)
        regularization = self.estimate_regularization(T.real)
        t = np.sqrt(T.clip(regularization))
        t_inv = np.diag(1.0 / t)
        v_inv = np.diag(1.0/np.sqrt(V.clip(1e-3)))
        self.sphere_ = np.dot(np.dot(np.dot(G, t_inv), G.T), v_inv)
        return self

    def transform(self, X, y=None, copy=None):
        """
        Parameters
        ----------
        X : array-like with shape [n_samples, n_features]
            The data to sphere along the features axis.
        """
        check_is_fitted(self, "mean_")
        X = as_float_array(X, copy=self.copy)
        return np.dot(X - self.mean_, self.sphere_.T)

def sphere_plate_zca_corr(plate):
    """
    sphere each plate to the DMSO negative control values
    Parameters:
    -----------
    plate: pandas.DataFrame
        dataframe of a single plate's featuredata and metadata
    Returns:
    -------
    pandas.DataFrame of the same shape as `plate`
    """
    # sphere featuredata to DMSO sphereing matrix
    sphereer = ZCA_corr()
    dmso_df = plate.loc[plate.Metadata_control_type=="negcon"]
    # dmso_df = plate.query("Metadata_pert_type == 'control'")
    dmso_vals = get_featuredata(dmso_df).to_numpy()
    all_vals = get_featuredata(plate).to_numpy()
    sphereer.fit(dmso_vals)
    sphereed_vals = sphereer.transform(all_vals)
    # concat with metadata columns
    feature_df = pd.DataFrame(
        sphereed_vals, columns=get_featurecols(plate), index=plate.index
    )
    metadata = get_metadata(plate)
    combined = pd.concat([feature_df, metadata], axis=1)
    assert combined.shape == plate.shape
    return combined

def calculate_percent_replicating_MOA(batch_path,plate,data_df=None):
    """
    For plates treated with the JUMP-MOA source plates, at least
    4 copies of each perturbation are present on each plate.
    Percent replicating is therefore calculated per plate.
    """
    metadata_compound_name = 'Metadata_pert_iname'
    n_samples_strong = 10000
    if type(data_df)!=pd.DataFrame:
        data_df = pd.read_csv(os.path.join(batch_path, plate,
                                            plate+'_normalized_feature_select_negcon_batch.csv.gz'))

    data_df = sphere_plate_zca_corr(data_df)

    data_df = remove_negcon_empty_wells(data_df)

    replicate_corr = list(corr_between_replicates(data_df, metadata_compound_name))
    null_corr = list(corr_between_non_replicates(data_df, n_samples=n_samples_strong, n_replicates=4, metadata_compound_name = metadata_compound_name))

    prop_95, value_95_replicating = percent_score(null_corr, replicate_corr)

    # return(prop_95)
    return replicate_corr, null_corr, prop_95, value_95_replicating

def calculate_percent_replicating_across_plates_MOA(batch_path1,plate1,batch_path2,plate2 ):
    """
    For plates treated with the JUMP-MOA source plates, at least
    4 copies of each perturbation are present on each plate.
    Percent replicating is therefore calculated per plate.
    """
    metadata_compound_name = 'Metadata_pert_iname'
    n_samples_strong = 10000

    data_df1 = pd.read_csv(os.path.join(batch_path1, plate1,
                                           plate1+'_normalized_feature_select_negcon.csv.gz'))
    data_df1 = sphere_plate_zca_corr(data_df1)
    data_df1 = remove_negcon_empty_wells(data_df1)

    data_df2 = pd.read_csv(os.path.join(batch_path2, plate2,
                                           plate2+'_normalized_feature_select_negcon.csv.gz'))
    data_df2 = sphere_plate_zca_corr(data_df2)
    data_df2 = remove_negcon_empty_wells(data_df2)

    replicate_corr = corr_between_replicates_across_plates(data_df1, data_df2)
    null_corr = corr_between_non_replicates_across_plates(data_df1, data_df2, n_samples=n_samples_strong)

    prop_95, _ = percent_score(null_corr, replicate_corr)

    return(prop_95)

def calculate_percent_matching_MOA(batch_path,plate, data_df=None):
    """
    For plates treated with the JUMP-MOA source plates, at least
    4 copies of each perturbation are present on each plate.
    Percent replicating is therefore calculated per plate.
    """
    metadata_moa_name = 'Metadata_moa'
    metadata_compound_name = 'Metadata_pert_iname'
    n_samples_strong = 10000

    if type(data_df)!=pd.DataFrame:
        data_df = pd.read_csv(os.path.join(batch_path, plate,
                                            plate+'_normalized_feature_select_negcon.csv.gz'))

    data_df = sphere_plate_zca_corr(data_df)

    data_df = remove_negcon_empty_wells(data_df)

    replicate_corr = list(corr_between_compound_moa(data_df, metadata_moa_name, metadata_compound_name))
    null_corr = list(null_corr_between_compound_moa(data_df, n_samples_strong, metadata_moa_name, metadata_compound_name))

    prop_95, value_95_matching = percent_score(null_corr, replicate_corr)

    return replicate_corr, null_corr, prop_95, value_95_matching

def calculate_percent_replicating_Target(batch_path,platelist,sphere=None,
suffix = '_normalized_feature_select_negcon.csv.gz',n_replicates=4):
    """
    For plates treated with the JUMP-Target source plates, most
    perturbations are only present in one or two 2 copies per plate.
    Percent replicating is therefore calculated per group of replicate plates.

    Since feature selection happens on a per-plate level, an inner join
    is performed across all plates in the replicate, meaning only common
    features are used in calculation of percent replicating.

    It doesn't look like sphering was done consistently in previous
    analysis of these plates, therefore it is configurable here; either
    not done, done at the plate level by passing 'sphere=plate', or
    done at the batch level by passing 'sphere=batch'.
    """
    metadata_compound_name = 'Metadata_broad_sample'
    n_samples_strong = 10000

    data_dict = {}

    for plate in platelist:
        plate_df = pd.read_csv(os.path.join(batch_path, plate,
                                            plate+suffix))

        if sphere == 'plate':
            plate_df = sphere_plate_zca_corr(plate_df)

        data_dict[plate] = plate_df

    data_df = pd.concat(data_dict, join='inner', ignore_index=True)

    if sphere == 'batch':
        data_df = sphere_plate_zca_corr(data_df)

    data_df = remove_negcon_empty_wells(data_df)

    replicate_corr = list(corr_between_replicates(data_df, metadata_compound_name))
    null_corr = list(corr_between_non_replicates(data_df, n_samples=n_samples_strong, n_replicates=n_replicates, metadata_compound_name = metadata_compound_name))

    prop_95, _ = percent_score(null_corr, replicate_corr)

    return(prop_95)

def calculate_percent_replicating_across_plates_Target(batch_path1,plate1,batch_path2,plate2 ):
    """
    For Target 1 vs Target 2
    """
    metadata_compound_name = 'Metadata_broad_sample'
    n_samples_strong = 10000

    data_df1 = pd.read_csv(os.path.join(batch_path1, plate1,
                                           plate1+'_normalized_feature_select_negcon.csv.gz'))
    data_df1 = remove_negcon_empty_wells(data_df1)

    data_df2 = pd.read_csv(os.path.join(batch_path2, plate2,
                                           plate2+'_normalized_feature_select_negcon.csv.gz'))
    data_df2 = remove_negcon_empty_wells(data_df2)

    replicate_corr = corr_between_replicates_across_plates(data_df1, data_df2,pertcol=metadata_compound_name)
    null_corr = corr_between_non_replicates_across_plates(data_df1, data_df2, n_samples=n_samples_strong,pertcol=metadata_compound_name)

    prop_95, _ = percent_score(null_corr, replicate_corr)

    return(prop_95)

def calculate_percent_matching_Target(batch_path_1,platelist_1,modality_1, batch_path_2,platelist_2, modality_2,
sphere=None,suffix = '_normalized_feature_select_negcon.csv.gz'):
    """

    It doesn't look like sphering was done consistently in previous
    analysis of these plates, therefore it is configurable here; either
    not done, done at the plate level by passing 'sphere=plate', or
    done at the batch level by passing 'sphere=batch'.
    """
    n_samples = 10000

    data_dict_1 = {}
    for plate in platelist_1:
        plate_df = pd.read_csv(os.path.join(batch_path_1, plate,
                                            plate+suffix))
        if sphere == 'plate':
            plate_df = sphere_plate_zca_corr(plate_df)

        data_dict_1[plate] = plate_df
    data_df_1 = pd.concat(data_dict_1, join='inner', ignore_index=True)
    if modality_1 =='Compounds':
        data_df_1.rename(columns={'Metadata_target':'Metadata_genes'},inplace=True)
    data_df_1['Metadata_modality'] = modality_1
    if sphere == 'batch':
        data_df_1 = sphere_plate_zca_corr(data_df_1)
    data_df_1 = remove_negcon_empty_wells(data_df_1)

    data_dict_2 = {}
    for plate in platelist_2:
        plate_df = pd.read_csv(os.path.join(batch_path_2, plate,
                                            plate+suffix))
        if sphere == 'plate':
            plate_df = sphere_plate_zca_corr(plate_df)

        data_dict_2[plate] = plate_df
    data_df_2 = pd.concat(data_dict_2, join='inner', ignore_index=True)
    if modality_2 =='Compounds':
        data_df_2.rename(columns={'Metadata_target':'Metadata_genes'},inplace=True)
    data_df_2['Metadata_modality'] = modality_2
    if sphere == 'batch':
        data_df_2 = sphere_plate_zca_corr(data_df_2)
    data_df_2 = remove_negcon_empty_wells(data_df_2)

    replicate_corr = list(correlation_between_modalities(data_df_1, data_df_2, modality_1, modality_2, 'Metadata_genes', 'Metadata_broad_sample'))
    null_corr = list(null_correlation_between_modalities(data_df_1, data_df_2, modality_1, modality_2, 'Metadata_genes', 'Metadata_broad_sample', n_samples))

    prop_95, _, _ = percent_score(null_corr, replicate_corr, how='both')

    return(prop_95)

def calculate_percent_matching_with_feature_dropout_Target(batch_path_1,platelist_1,modality_1, batch_path_2,platelist_2, modality_2,
sphere=None,suffix = '_normalized_negcon.csv.gz',drop='AGP'):
    """

    It doesn't look like sphering was done consistently in previous
    analysis of these plates, therefore it is configurable here; either
    not done, done at the plate level by passing 'sphere=plate', or
    done at the batch level by passing 'sphere=batch'.
    """
    import pycytominer
    n_samples = 10000

    data_dict_1 = {}
    for plate in platelist_1:
        plate_df = pd.read_csv(os.path.join(batch_path_1, plate,
                                            plate+suffix))
        cols_to_drop = [x for x in plate_df.columns if drop in x]
        plate_df.drop(columns=cols_to_drop,inplace=True)
        feature_select_features = pycytominer.cyto_utils.infer_cp_features(
        plate_df
        )
        plate_df = pycytominer.feature_select(
        profiles=plate_df,
        features=feature_select_features,
        operation=['variance_threshold','correlation_threshold',
        'drop_na_columns','blocklist']
        )
        if sphere == 'plate':
            plate_df = sphere_plate_zca_corr(plate_df)

        data_dict_1[plate] = plate_df
    data_df_1 = pd.concat(data_dict_1, join='inner', ignore_index=True)
    if modality_1 =='Compounds':
        data_df_1.rename(columns={'Metadata_target':'Metadata_genes'},inplace=True)
    data_df_1['Metadata_modality'] = modality_1
    if sphere == 'batch':
        data_df_1 = sphere_plate_zca_corr(data_df_1)
    data_df_1 = remove_negcon_empty_wells(data_df_1)

    data_dict_2 = {}
    for plate in platelist_2:
        plate_df = pd.read_csv(os.path.join(batch_path_1, plate,
                                            plate+suffix))
        cols_to_drop = [x for x in plate_df.columns if drop in x]
        plate_df.drop(columns=cols_to_drop,inplace=True)
        feature_select_features = pycytominer.cyto_utils.infer_cp_features(
        plate_df
        )
        plate_df = pycytominer.feature_select(
        profiles=plate_df,
        features=feature_select_features,
        operation=['variance_threshold','correlation_threshold',
        'drop_na_columns','blocklist']
        )
        if sphere == 'plate':
            plate_df = sphere_plate_zca_corr(plate_df)

        data_dict_2[plate] = plate_df
    data_df_2 = pd.concat(data_dict_2, join='inner', ignore_index=True)
    if modality_2 =='Compounds':
        data_df_2.rename(columns={'Metadata_target':'Metadata_genes'},inplace=True)
    data_df_2['Metadata_modality'] = modality_2
    if sphere == 'batch':
        data_df_2 = sphere_plate_zca_corr(data_df_2)
    data_df_2 = remove_negcon_empty_wells(data_df_2)

    replicate_corr = list(correlation_between_modalities(data_df_1, data_df_2, modality_1, modality_2, 'Metadata_genes', 'Metadata_broad_sample'))
    null_corr = list(null_correlation_between_modalities(data_df_1, data_df_2, modality_1, modality_2, 'Metadata_genes', 'Metadata_broad_sample', n_samples))

    prop_95, _, _ = percent_score(null_corr, replicate_corr, how='both')

    return(prop_95)

def plot_simple_comparison(df,x,hue,y='Percent Replicating',order=None,hue_order=None,
col=None, col_order=None, col_wrap=None,row=None,row_order=None,jitter=0.25,dodge=True,plotname=None,
ylim=None, title=None,aspect=1,sharex=True,facet_kws={}):
    plt.rcParams["legend.markerscale"] =1.5
    sns.set_style("ticks")
    sns.set_context("paper",font_scale=1.5)
    g = sns.catplot(data=df, x = x ,y = y, order=order,
    hue=hue, hue_order=hue_order, col=col, col_order = col_order, col_wrap=col_wrap,row=row,
    row_order = row_order, palette='Set1',s=12,linewidth=1,jitter=jitter,
    alpha=0.8,dodge=dodge,aspect=aspect,sharex=sharex,facet_kws=facet_kws)
    if sharex:
        labels = []
        if not order:
            orig_labels = list(dict.fromkeys(df[x].values).keys())
        else:
            orig_labels = order
        for label in orig_labels:
            if type(label)!= str:
                label = str(int(label))
            labels.append(textwrap.fill(label, width=45/len(orig_labels),break_long_words=False))
        g.set_xticklabels(labels=labels,rotation=0)
    if ylim:
        ymin,ymax=ylim
    else:
        ymin = 50
        ymax = 80
    if df[y].min()<ymin:
        ymin = df[y].min()-2
    if df[y].max()>ymax:
        ymax = df[y].max()+2
    g.set(ylim=([ymin,ymax]))
    if plotname:
        plotname = f"../figures/{plotname}"
    else:
        plotname = f"../figures/{x}-{y}-{hue}-{col}-{row}.png"
    if not col:
        if not row:
            if title:
                g.set(title=title)
            else:
                g.set(title=f"{x}-{y}")
    g.savefig(plotname,dpi=300)
    print(f'Saved to {plotname}')

def plot_two_comparisons(df,x='Percent Replicating',y='Percent Matching',hue = None, hue_order=None,
col=None, col_order=None,col_wrap=None,row=None,row_order=None,style=None,xlim=None,ylim=None,title=None,
title_variable = None, facet_kws={'sharex':True}):
    plt.rcParams["legend.markerscale"] =1.5
    sns.set_style("ticks")
    sns.set_context("paper",font_scale=1.5)
    g = sns.relplot(data=df, x = x ,y= y, hue=hue, hue_order=hue_order, col=col, col_order = col_order,
    col_wrap=col_wrap, row=row, row_order = row_order, style = style, palette='Set1',edgecolor='k',alpha=0.8,s=80,
    facet_kws=facet_kws)
    if xlim:
        xmin,xmax=xlim
    else:
        xmin = 50
        xmax = 80
    if df[x].min()<xmin:
        xmin = df[x].min()-2
    if df[x].max()>xmax:
        xmax = df[x].max()+2
    g.set(xlim=([xmin,xmax]))
    if ylim:
        ymin,ymax=ylim
    else:
        ymin = 5
        ymax = 40
    if df[y].min()<ymin:
        ymin = df[y].min()-2
    if df[y].max()>ymax:
        ymax = df[y].max()+2
    g.set(ylim=([ymin,ymax]))
    if title:
        g.set(title=title)
    elif title_variable:
        g.set(title=title_variable)
    plotname = f"../figures/{x}-{y}-{hue}-{col}-{row}-{style}.png"
    g.savefig(plotname,dpi=300)
    print(f'Saved to {plotname}')

def enforce_modality_match_order(modality1,modality2):
    modality_dict_forward = {'Compounds':1,'ORF':2,'CRISPR':3}
    modality_dict_reverse = {v:k for k,v in modality_dict_forward.items()}
    modality_list = [modality_dict_forward[modality1],modality_dict_forward[modality2]]
    modality_list.sort()
    return f"{modality_dict_reverse[modality_list[0]]} - {modality_dict_reverse[modality_list[1]]}"

def enforce_timepoint_order(timepoint1,timepoint2):
    timepoint_list = [int(timepoint1),int(timepoint2)]
    timepoint_list.sort()
    return f"{timepoint_list[0]}-{timepoint_list[1]}"

def enforce_timepoint_order_in_plot(timepointlist):
    timepointlist=list(set(timepointlist))
    intlist = []
    timepoint_dict = {}
    for x in timepointlist:
        first,second=x.split('-')
        intlist.append([int(first),int(second)])
    intlist.sort()
    outlist = []
    for x in intlist:
        outlist.append(f"{x[0]}-{x[1]}")
    return outlist

def enforce_modality_match_order_in_plot(modalitylist):
    modalitylist=list(set(modalitylist))
    modality_dict_forward = {'Compounds':1,'ORF':2,'CRISPR':3}
    modality_dict_reverse = {v:k for k,v in modality_dict_forward.items()}
    master_modalitylist = []
    for eachpair in modalitylist:
        modality1,modality2 = eachpair.split(' - ')
        sublist = [modality_dict_forward[modality1],modality_dict_forward[modality2]]
        sublist.sort()
        master_modalitylist.append(sublist)
    master_modalitylist.sort()
    outlist = [f"{modality_dict_reverse[x[0]]} - {modality_dict_reverse[x[1]]}" for x in master_modalitylist]
    return outlist

def safe_literal_eval(node):
    """
    If eval doesn't work, make it a np.nan
    For JUMP-Scope, this is the only other string
    found in lists
    """
    try:
        return ast.literal_eval(node)
    except SyntaxError:
        return np.nan
    except ValueError:
        return np.nan

def jitter_data(data, magnitude=0.2):
    noise = np.random.uniform(low=-magnitude, high=magnitude, size=data.shape)
    return data + noise

def group_plot(
    df,
    x,
    y,
    group,
    error_x=None,
    error_y=None,
    fig=None,
    ax_=None,
    legend=False,
    legend_title=None,
    legend_location=None,
    infer_custom_legend=False,
    s=None,
    jitter_x=None,
    use_markers=False,
    label=None,
    alpha=None,
    x_lim=None,
    y_lim=None,
    plot_title=False,
    xlabel=None,
    ylabel=None
    ):
    cmap = plt.cm.tab10
    colour_palette = list()
    markers = "sxo^D2"
    for i in range(cmap.N):
        colour_palette.append(cmap(i))

    if fig is None and ax_ is None:
        fig, ax = plt.subplots()
    else:
        ax = ax_
    for i, (group_label, group_df) in enumerate(df.groupby(group)):
        if infer_custom_legend:
            if len(group_df) > 1:
                raise ValueError("Cannot infer custom legend from grouped DF of more than 1")
            # Columns to drop that are not to be used for paired comparison
            infer_legend = df.apply(pd.Series.duplicated).any().drop([
                "Batch",
                "Assay_Plate_Barcode",
                "Percent_Replicating",
                "Percent_Matching",
                "euclidean_distance",
                "cell_count"])
            infer_legend = infer_legend[~infer_legend].index
        if error_x is not None or error_y is not None:
            ax.errorbar(
                group_df.loc[:, x],
                group_df.loc[:, y],
                xerr=group_df[error_x] if error_x else None,
                yerr=group_df[error_y] if error_y else None,
                ecolor=colour_palette[i],
                # label=group_label,
                # fmt="None",
                # zorder=1,
                lw=2
            )
        else:
            ax.scatter(
                group_df.loc[:, x] if not jitter_x else jitter_data(group_df.loc[:, x].values, jitter_x),
                group_df.loc[:, y],
                color=colour_palette[i],
                # We use .iloc[0] here becasue we assume that every group has a len(df) of 1
                # Labels cannot be resolved for groups larger than 1
                label=group_label if not infer_custom_legend else group_df[infer_legend].iloc[0].to_dict(),
                alpha=alpha,
                marker=markers[i] if use_markers else None,
                s=s
                # zorder=2
            )
        if jitter_x:
            # A solution I hate but it works for the specific case of figure S1C.
            # Likely to be a nuisance if jitter is added to any other type of plot
            ax.set_xticks(np.arange(df.loc[:, x].min(), df.loc[:, x].max() + 1))

        if x_lim:
            ax.set_xlim(x_lim)

        if y_lim:
            ax.set_ylim(y_lim)

        if legend:
            ax.legend(loc=legend_location if legend_location else None, title=legend_title)

        if label:
            for col, rows in label.items():
                for item in rows:
                    annotate_x = group_df[group_df[col] == item][x].values
                    annotate_y = group_df[group_df[col] == item][y].values
                    ax.annotate(item, annotate_x, annotate_y)

    if plot_title:
        ax.set_title(group if not plot_title else plot_title, size=15)
    ax.set_xlabel(x if not xlabel else xlabel, fontsize=15)
    ax.set_ylabel(y if not ylabel else ylabel, fontsize=15)
    plt.tight_layout()
    fig.set_facecolor("white")
    plt.subplots_adjust(
                    wspace=0.2,
                    hspace=0.2)

def find_group_avg_df(_df, group, **kwargs):
    df = _df.copy()
    # Define operations to perform on columns during aggregation
    agg_dict = {
        'Percent_Replicating' : lambda x: list(x),
        'Mean_Percent_Replicating' : lambda y: np.mean(y),
        'SD_Percent_Replicating' : lambda z: float('%.3f'%np.std(z)),
        'Percent_Matching' : lambda x: list(x),
        'Mean_Percent_Matching' : lambda y: np.mean(y),
        'SD_Percent_Matching' : lambda z: float('%.3f'%np.std(z)),
    }
    # Make new columns
    df['Mean_Percent_Replicating'] = list(df['Percent_Replicating'])
    df['SD_Percent_Replicating'] = list(df['Percent_Replicating'])
    df['Mean_Percent_Matching'] = list(df['Percent_Matching'])
    df['SD_Percent_Matching'] = list(df['Percent_Matching'])
    # Implement desired cols and define operations
    if "add_cols" in kwargs:
        for new_col in kwargs["add_cols"].values():
            df[new_col[1]] = list(df[new_col[0]])
            # new_col[2] is the function to apply
            agg_dict.update({new_col[1]: new_col[2]})
    # Perform aggregation
    group_df = df.groupby(group,as_index=False).agg(agg_dict)
    return group_df

def profile_finder(df, profile_path, profile_type):
    """Returns a list of DataFrames from the provided df"""
    # Store profile types
    type_dict = {
        "unnormalized": ".csv.gz",
        "norm_feature_selected": "_normalized_feature_select_negcon_plate.csv.gz"
        }

    # Select desired profile type
    suffix = type_dict[profile_type]

    output = list()
    plate_info = list()

    for ind, row in df.iterrows():
        df_path = os.path.join(profile_path, row["Batch"], row["Assay_Plate_Barcode"], f"{row['Assay_Plate_Barcode']}{suffix}")
        print(f"loading: {df_path}")
        output.append(pd.read_csv(df_path))
        plate_info.append({"Batch": row["Batch"], "Assay_Plate_Barcode": row["Assay_Plate_Barcode"]})

    return output, plate_info

def compare_paired_correlations(
    df: pd.DataFrame,
    profile_path: str,
    correlation_metric: str,
    grouping_metric: str,
    fig=None,
    ax=None,
    plot_title=None,
    xlabel=None,
    ylabel=None,
    ):
    """
    Between two profiles, return a dataframe that contains only the desired feature
    """
    # Load requested profiles
    profiles, plate_info = profile_finder(df, profile_path, "unnormalized")

    # Keep only the correlation columns
    for i, i_df in enumerate(profiles):
        mask_cols = [] # Store the columns to mask, find for each profile independently
        for cols in i_df.columns:
            if correlation_metric in cols and "BrightField" not in cols:
                mask_cols.append(cols)
        # Mask columns
        profiles[i] = i_df[mask_cols]

        unique_value = df[
            (df["Batch"].str.contains(plate_info[i]["Batch"])) &
            (df["Assay_Plate_Barcode"].str.contains(plate_info[i]["Assay_Plate_Barcode"]))
            ][grouping_metric].values

        profiles[i][grouping_metric] = unique_value[0]

    profiles = pd.concat(profiles)

    for j in profiles.groupby(grouping_metric):
        print(j[1].shape[0])

    profiles = profiles.melt(id_vars=grouping_metric)

    fig.set_facecolor("white")

    seaborn.violinplot(
        ax=ax,
        data=profiles,
        y="variable",
        x="value",
        orient="h",
        hue=grouping_metric,
        split=True,
        cut=0,
        figsize=(20, 10),
        palette="muted"
        )

    ax.set_title(None if not plot_title else plot_title, size=15)
    ax.set_xlabel(None if not xlabel else xlabel, fontsize=15)
    ax.set_ylabel(None if not ylabel else ylabel, fontsize=15)
    plt.tight_layout()

def compare_single_correlations(
    df: pd.DataFrame,
    profile_path: str,
    correlation_metric: str,
    fig=None,
    ax=None,
    plot_title=None,
    xlabel=None,
    ylabel=None,
    ):
    """
    Between two profiles, return a dataframe that contains only the desired feature
    """
    # Load requested profiles
    profiles, _ = profile_finder(df, profile_path, "unnormalized")
    # Keep only the correlation columns
    for i, i_df in enumerate(profiles):
        mask_cols = [] # Store the columns to mask, find for each profile independently
        for cols in i_df.columns:
            if correlation_metric in cols and "BrightField" not in cols:
                mask_cols.append(cols)
        # Mask columns
        profiles[i] = i_df[mask_cols]

    profiles = pd.concat(profiles)

    profiles = profiles.melt()

    fig.set_facecolor("white")

    seaborn.boxenplot(
        ax=ax,
        data=profiles,
        y="variable",
        x="value",
        orient="h",
        showfliers=False
        )

    ax.set_title(None if not plot_title else plot_title, size=15)
    ax.set_xlabel(None if not xlabel else xlabel, fontsize=15)
    ax.set_ylabel(None if not ylabel else ylabel, fontsize=15)
    plt.tight_layout()

def format_column_names(column_name):
    """
    For a given string, hyphens and underscores are replaced with spaces and
    first letters are capitalised.
    """
    return column_name.title().replace("-", " ").replace("_", " ")

def aggregate_duplicates(df, non_group_cols):
    """
    Group data on setting columns, calculate the mean for grouped rows
    (which are therefore duplicates).

    The returned dataframe will not contain Assay_Plate_Barcode, since
    the unique barcodes cannot be reconciled into an aggregated mean,
    so are therefore dropped.
    """
    # Columns to not be used for grouping
    # non_group_cols = [
    #     "Assay_Plate_Barcode",
    #     "Batch",
    #     "Vendor",
    #     "value_95_replicating",
    #     "Percent_Replicating",
    #     "Size_MB",
    #     "Size_MB_std",
    #     "Percent_Matching",
    #     "value_95_matching",
    #     "cell_count",
    #     "Sites-SubSampled",
    #     "BF_Zplanes", # Ignore since only one BF zplane is used
    #     "brightfield_z_plane_used",
    #     # Remove channel names due to some profiles having "AGP" features
    #     # and others having "WGPhalloidin" instead
    #     "feature_channels_found",
    #     "channel_names"
    # ]
    # Find the columns that are not in non_group_cols
    diff = list(set(df.columns) - set(non_group_cols))
    # Group df by setting columns, find the mean, then reset the index
    df = df.groupby(diff, dropna=False, as_index=False).mean()
    return df

def make_leaderboard(df, columns, non_group_cols, average_duplicates=True):
    """
    Process match_rep_df into a nice leaderboard
    """
    if average_duplicates:
        df = aggregate_duplicates(df, non_group_cols)

    # Create an aggregation score for replicating/matching and normalize to the max value
    df["Percent_Score"] = df[["Percent_Replicating", "Percent_Matching"]].mean(axis=1)
    df["Percent_Score"] = (df["Percent_Score"] / df["Percent_Score"].max()) * 100
    df = df.round({"Percent_Score": 1})

    # Adjust % rep/match to that of the max
    df["Percent_Replicating"] = (df["Percent_Replicating"] / df["Percent_Replicating"].max()) * 100
    df["Percent_Matching"] = (df["Percent_Matching"] / df["Percent_Matching"].max()) * 100
    df = df.round({"Percent_Replicating": 1})
    df = df.round({"Percent_Matching": 1})

    # Sort based on the Percent_Score and then add Place (ie 1st, 2nd, 3rd etc.)
    df = df.sort_values("Percent_Score", ascending=False)[columns]
    df["Place"] = df.reset_index(drop=True).index+1

    # Move place column to first position
    df.insert(0, "Place", df.pop("Place"))

    # Format column_names
    df.rename(columns=format_column_names, inplace=True)
    df.rename(columns={"Z Plane": "Z Planes"}, inplace=True)
    df.rename(columns={"Aperture": "NA"}, inplace=True)
    df.rename(columns={"Dry Immersion": "Immersion"}, inplace=True)
    return df

def aggregate_comparison_pvalues(df: pd.DataFrame, remove_cols: list, aggregate_rows: bool, aggregate_similar_features: bool):
    # Remove columns that are in remove cols and also only keep columns that are pvalues
    df = df.loc[:, (~df.columns.str.startswith(tuple(remove_cols))) & (df.columns.str.contains("_pvalue"))]
    # Drop Location, Number, Parent and Neighbour features
    df = df.loc[:, ~(df.columns.str.contains("|".join(["Location", "Number", "Parent", "Neighbors", "Children"])))]
    if aggregate_similar_features:
        compartments = ["Nuclei", "Cytoplasm", "Cells"]
        result_df = pd.DataFrame()

        for compart in compartments:
            subset_cols = [col for col in df.columns if compart in col]
            sub_df = df[subset_cols]
            # print(sub_df.columns)

            # Find the first 3 words of column names that will be used for grouping
            group_words = [col.split("_")[:3] for col in subset_cols]
            # Join them back
            group_words = ["_".join(w) for w in group_words]

            # Find the mean for the feature groups
            # result_df = pd.concat([result_df, sub_df.groupby(group_words, axis=1).mean().reset_index()])
            result_df = result_df.join(sub_df.groupby(group_words, axis=1).mean(), how="right")

        df = result_df
    # Convert DF to long format
    if aggregate_rows:
        df = df.mean(axis=0).to_frame().T.melt(var_name="feature", value_name="pvalue")
    else:
        df = df.melt(var_name="feature", value_name="pvalue")
    df["feature_type"] = df["feature"].str.split("_").str[1]
    return df

def combine_evalzoo_metrics(match_rep_df: pd.DataFrame, evalzoo_dir: str = "results/results"):
    # Add Batch and Metadata_Plate since we need to check the presence of a
    # string in them later, and you can't check them if they don't exist
    output_df = pd.DataFrame({"Batch": [], "Metadata_Plate": []})
    for ind, row in match_rep_df.iterrows():
        # Need to do some check that this row has not already been processed (so
        # we don't add it twice)
        if not (row["Batch"] in output_df["Batch"].values and row["Assay_Plate_Barcode"] in output_df["Metadata_Plate"].values):
            ref_metric_path = os.path.join(evalzoo_dir, row["Batch"]+"_ref", f"metrics_level_1_ref.parquet")
            non_rep_metric_path = os.path.join(evalzoo_dir, row["Batch"]+"_non_rep", f"metrics_level_1_non_rep.parquet")

            ref_metric_df = pd.read_parquet(ref_metric_path)
            non_rep_metric_df = pd.read_parquet(non_rep_metric_path)

            # Since we are going to use the metric_type to differentiate between ref
            # and non_rep, remove this information from the feature names
            # ref_metric_df = ref_metric_df.rename(lambda x: x + "_ref" if "sim_" in x else x, axis=1)
            # non_rep_metric_df = non_rep_metric_df.rename(lambda x: x + "_non_rep" if "sim_" in x else x, axis=1)

            ref_metric_df = ref_metric_df.rename(lambda x: x.replace("_ref", "") if "sim_" in x else x, axis=1)
            non_rep_metric_df = non_rep_metric_df.rename(lambda x: x.replace("_non_rep", "") if "sim_" in x else x, axis=1)

            # Add cols
            (
                ref_metric_df["Vendor"],
                ref_metric_df["Batch"],
                ref_metric_df["Magnification"],
                ref_metric_df["Binning"],
                ref_metric_df["z_plane"],
                ref_metric_df["metric_type"]
                ) = (
                row["Vendor"],
                row["Batch"],
                row["Magnification"],
                row["Binning"],
                row["z_plane"],
                "ref"
            )

            (
                non_rep_metric_df["Vendor"],
                non_rep_metric_df["Batch"],
                non_rep_metric_df["Magnification"],
                non_rep_metric_df["Binning"],
                non_rep_metric_df["z_plane"],
                non_rep_metric_df["metric_type"]
                ) = (
                row["Vendor"],
                row["Batch"],
                row["Magnification"],
                row["Binning"],
                row["z_plane"],
                "non_rep"
                )

            # Cast some columns to int
            ref_metric_df["Magnification"] = ref_metric_df["Magnification"].astype("Int64")
            ref_metric_df["Binning"] = ref_metric_df["Binning"].astype("Int64")
            ref_metric_df["z_plane"] = ref_metric_df["z_plane"].astype("Int64")

            non_rep_metric_df["Magnification"] = non_rep_metric_df["Magnification"].astype("Int64")
            non_rep_metric_df["Binning"] = non_rep_metric_df["Binning"].astype("Int64")
            non_rep_metric_df["z_plane"] = non_rep_metric_df["z_plane"].astype("Int64")

            output_df = pd.concat([output_df, ref_metric_df, non_rep_metric_df], axis=0)

    return output_df

def download_from_cpg(prefix: str, save_dir: str):
    """Download all images that match the prefix.
    Prefix must be the complete path to a particular file or
    group of files.

    The prefix, in essence, works like a wildcard search. If a
    directory contains file1, file2 and file3 and the prefix is
    'file', all files will be downloaded. If the prefix is file1,
    only file1 will be downloaded."""

    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)

    s3 = boto3.resource('s3')
    # The Cell Painting Gallery allows for unsigned requests
    s3.meta.client.meta.events.register('choose-signer.s3.*', disable_signing)
    # Define the bucket
    bucket = s3.Bucket("cellpainting-gallery")
    for obj in bucket.objects.filter(Prefix=prefix):
        print(f"Downloading: {obj.key}")
        s3.Bucket("cellpainting-gallery").download_file(obj.key, os.path.join(save_dir, obj.key.split('/')[-1]))

def plot_images(img_paths: list, crop: tuple = None, x0_y0: tuple = (0, 0), plot_title: str = None, ch_names:dict = None):
    if ch_names is None:
        ch_names = {
            1: "DNA",
            2: "Endoplasmic reticulum",
            3: "Nucleoli, cytoplasmic RNA",
            4: "Actin, Golgi, plasma membrane",
            5: "Mitochondria"
        }
    fig, ax = plt.subplots(1, len(img_paths), figsize=(20, 5*len(img_paths)))
    for i, path in enumerate(img_paths):
        img = skimage.io.imread(path)
        if crop:
            img = img[x0_y0[1]:crop[1]+x0_y0[1], x0_y0[0]:crop[0]+x0_y0[0]]
        minq = np.quantile(img, 0.02)
        maxq = np.quantile(img, 0.98)
        img = np.clip(img, minq, maxq)
        ax[i].imshow(img, cmap="gray")
        title = path[-7:-4]
        ax[i].set_title(ch_names.get(i+1), fontsize=17)
        ax[i].axis('off')
    if plot_title is not None:
        fig.set_title(plot_title)
    fig.tight_layout()
    fig.set_facecolor("white")

    return fig

def plot_two_distributions(
        df: pd.DataFrame,
        x: str,
        y: str,
        add_stat_annotation: bool = True,
        xlabel: list = None,
        ylabel: str = None,
        plot_title: str = None,
        include_pval: bool = False,
        ax = None,
        fig = None
):
    assert len(df[x].unique()) == 2, f"Can't compare anything other than two distributions. Got {df[x].unique()} for x"
    group1, group2 = sorted(df[x].unique())

    if ax is None or fig is None:
        fig, ax = plt.subplots(figsize=(5,5))
        ax.tick_params(labelsize=12)

    sns.boxenplot(data=df, x=x, y=y, ax=ax, showfliers=False)
    sns.stripplot(data=df, x=x, y=y, ax=ax, color="black")

    if add_stat_annotation:
        # Define where to draw the stat annotation
        left_dist, right_dist, y_max, height_addition = 0, 1, df[y].max() + 0.1, 0.1

        _, pval = scipy.stats.ttest_ind(df[df[x]==group1][y].values, df[df[x]==group2][y].values)
        pval = round(pval, 3)
        print(f"The pvalue is: {pval}")
        if include_pval:
            significance = "ns" if pval > 0.05 else f"pvalue: {pval}"
            if pval > 0.05:
                significance = significance + f"\npvalue: {pval}"
        else:
            significance = "ns" if pval > 0.05 else f"pvalue: {pval}"

        # Give some extra room to the annotation
        ax.set_ylim(0, y_max*1.25)

        # Draw the annotation
        ax.plot([left_dist, left_dist, right_dist, right_dist], [y_max, y_max+height_addition, y_max+height_addition, y_max], lw=1, c="black")
        ax.text((left_dist+right_dist)*.5, y_max+height_addition, significance, ha="center", va="bottom", color="black", fontsize=12)

        tstat, pval = scipy.stats.ttest_ind(df[df[x]==group1][y].values, df[df[x]==group2][y].values)
        print(f"The pvalue is: {pval}")

    ax.set_xlabel(x if not xlabel else xlabel, fontsize=15)
    ax.set_ylabel(y if not ylabel else ylabel, fontsize=15)
    ax.set_title(plot_title if plot_title is not None else None)
    plt.tight_layout()
    fig.set_facecolor("white")
    plt.subplots_adjust(
                    wspace=0.2,
                    hspace=0.2)
    return fig

def plot_feature_density(
        df: pd.DataFrame,
        x: str,
        group: str,
        fig = None,
        ax = None,
        xlabel: list = None,
        ylabel: str = None,
        plot_title: str = None,
        ):

    if fig is None or ax is None:
        fig, ax = plt.subplots(figsize=(5, 5))

    kde = sns.kdeplot(
        data=df,
        x=x,
        hue=group,
        ax=ax,
        # bw_method=0.2,
        # bw_method="silverman"
        # log_scale=[False, True]
        )

    # Normalize feature groups to between 0-1
    for line in kde.lines:
        y = line.get_data()[1]
        line.set_data(line.get_data()[0], y / y.max())

    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1.1)
    ax.set_xlabel(x if not xlabel else xlabel, fontsize=15)
    ax.set_ylabel(y if not ylabel else ylabel, fontsize=15)
    if plot_title is not None:
        ax.set_title(plot_title, fontsize=18)
    plt.tight_layout()
    fig.set_facecolor("white")
    return fig


import sys
import os

class Pen:
  def __init__(self):
    self.x=0
    self.y=0
    self.down=True

  def moveto(self,dx,dy,fp):
    newx=self.x+dx
    newy=self.y+dy
    if self.down:
      print('  <line x1="%f" y1="%f" x2="%f" y2="%f" stroke="black" style="stroke-width: 2;" />'%(self.x,self.y,newx,newy),file=fp)
    self.x=newx
    self.y=newy

  def circle(self,r,fp):
    if self.down:
      print('  <ellipse cx="%f" cy="%f" rx="%f" ry="%f" stroke="black" style="stroke-width: 2;" />'%(self.x,self.y,r,r),file=fp)

  def reset(self):
    self.x=0
    self.y=0
    self.down=True

if len(sys.argv)!=2:
  print('usage: python typo2svg.py <filename.typo>')
  exit()

ftypo=sys.argv[1]
svgdir=ftypo.split('.')[0]
p=Pen()
print('creating dir: %s'%svgdir)
os.mkdir(svgdir)

with open(ftypo) as ft:
  cmaplns=ft.readlines()

for ln in cmaplns:
  parts=ln.split('=>')
  c=parts[0]
  if c==' ':
    continue
  ilist=parts[1].rstrip().split(';')
  print('[codegen] working on file: %s/%s.svg'%(svgdir,c))
  p.reset()
  with open('%s/%s.svg'%(svgdir,c),'w') as fsvg:
    print('<svg viewBox="0 0 100 100" xmlns="http://www.w3.org/2000/svg">',file=fsvg)
    for inst in ilist:
      if inst=='u':
        p.down=False
      elif inst=='d':
        p.down=True
      elif inst[0]=='c':
        cr=inst.split(' ')
        p.circle(float(cr[1]),fsvg)
      else:
        xy=inst.split(' ')
        p.moveto(float(xy[0]),float(xy[1]),fsvg)
    print('</svg>',file=fsvg)