ED_LRR/ed_lrr_gui/gui/main.py

685 lines
24 KiB
Python

# -*- coding: utf-8 -*-
import csv
import gzip
import multiprocessing as MP
import os
import pathlib
import queue
import sys
from sys import exit
from datetime import datetime, timedelta
from urllib.request import Request, urlopen
import _ed_lrr
import ed_lrr_gui
import requests as RQ
from ed_lrr_gui import Preprocessor, Router, cfg
from ed_lrr_gui.gui.ed_lrr import Ui_ED_LRR
from ed_lrr_gui.gui.widget_route import Ui_diag_route
from PyQt5.QtCore import QObject, Qt, QThread, QTimer, pyqtSignal
from PyQt5.QtGui import QColor, QPalette, QIcon
from PyQt5.QtWidgets import (
QAction,
QApplication,
QFileDialog,
QMainWindow,
QMessageBox,
QProgressDialog,
QTreeWidgetItem,
QLabel,
QDialog,
)
class ProgressDialog(QProgressDialog):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setWindowModality(Qt.WindowModal)
def sizeof_fmt(num, suffix="B"):
for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
if abs(num) < 1024.0:
return "{:.02f} {}{}".format(num, unit, suffix)
num /= 1024.0
return "{:.02f} {}{}".format(num, "Yi", suffix)
def t_round(dt):
return dt - dt % timedelta(seconds=1)
class Job(QObject):
progress = pyqtSignal("PyQt_PyObject")
def __init__(self, app, main_window, cls, *args, **kwargs):
super().__init__()
self.job = cls(*args, **kwargs)
self.timer = QTimer(app)
self.app = app
self.main_window = main_window
self.timer.timeout.connect(self.interval)
self.timer.start(100)
self.last_val = None
self.progress_dialog = None
self.state = {}
self.setup_progress(self.handle_progess)
def setup_progress(self, handle_progess):
self.progress.connect(lambda *args, **kwargs: handle_progess(*args, **kwargs))
def handle_progess(self, *args, **kwargs):
raise NotImplementedError
def start(self):
if self.progress_dialog is None:
self.progress.connect(
lambda *args, **kwargs: print("PROGRESS:", *args, **kwargs)
)
self.started = datetime.today()
return self.job.start()
def cancel(self):
self.job.terminate()
self.job = None
def __bool__(self):
return not self.done
@property
def done(self):
if self.job:
return (self.job.is_alive() == False) and (self.job.queue.empty())
return True
def interval(self):
while self:
try:
res = self.job.queue.get(True, 0.1)
except queue.Empty:
return
if res == self.last_val:
continue
self.state.update(res)
self.progress.emit(self.state)
self.last_val = res
class PreprocessJob(Job):
def __init__(self, app, main_window, *args, **kwargs):
super().__init__(app, main_window, Preprocessor, *args, **kwargs)
self.progress_dialog = ProgressDialog("", "Cancel", 0, 0, self.main_window)
self.progress_dialog.setAutoClose(False)
self.progress_dialog.canceled.connect(self.cancel)
self.progress_dialog.show()
self.start()
def handle_progess(self, state):
sent = object()
if state.get("return", sent) != sent:
self.progress_dialog.close()
return
msg = "Processed: {}/{}".format(
sizeof_fmt(state["status"]["done"]), sizeof_fmt(state["status"]["total"])
)
state["status"]["prc_done"] = (
state["status"]["done"] / state["status"]["total"]
) * 100
title = "[{prc_done:.2f}%] Processing {file}".format(**state["status"])
self.progress_dialog.setMinimum(0)
self.progress_dialog.setMaximum(100 * 100)
self.progress_dialog.setWindowTitle(title)
self.progress_dialog.setLabelText(msg)
self.progress_dialog.setValue(int(state["status"]["prc_done"] * 100))
class RouterJob(Job):
def __init__(self, app, main_window, *args, **kwargs):
super().__init__(app, main_window, Router, *args, **kwargs)
self.progress_dialog = ProgressDialog("", "Cancel", 0, 0, self.main_window)
self.progress_dialog.setAutoClose(False)
self.progress_dialog.setLabelText("Loading data (this will take a bit) ...")
self.progress_dialog.setWindowTitle("Loading...")
self.progress_dialog.canceled.connect(self.cancel)
self.progress_dialog.show()
self.start()
self.state = {}
def handle_progess(self, state):
sent = object()
if state.get("return", sent) != sent:
print(state["return"])
self.progress_dialog.close()
route_win = WRoute(self.main_window, state["return"])
return
msg = "Depth: {depth}\nBody: {body}\nQueued: {queue_size}\nDistance remaining: {d_rem:.2f} Ly".format(
**state["status"]
)
title = "[{prc_done:.2f}%] Plotting route from [{from}] to [{to}]".format(
**state["status"]
)
self.progress_dialog.setMinimum(0)
self.progress_dialog.setMaximum(100 * 100)
self.progress_dialog.setWindowTitle(title)
self.progress_dialog.setLabelText(msg)
self.progress_dialog.setValue(int(state["status"]["prc_done"] * 100))
class DownloadThread(QThread):
progress = pyqtSignal("PyQt_PyObject")
def __init__(self, systems_url, systems_file, bodies_url, bodies_file):
super().__init__()
self.systems_url = systems_url
self.systems_file = systems_file
self.bodies_url = bodies_url
self.bodies_file = bodies_file
self.running = True
def __del__(self):
self.wait()
def stop(self):
self.running = False
def run(self):
dl_jobs = [
(self.systems_url, self.systems_file),
(self.bodies_url, self.bodies_file),
]
for url, dest in dl_jobs:
outfile = url.split("/")[-1]
size = RQ.head(url, headers={"Accept-Encoding": "None"})
size.raise_for_status()
size = int(size.headers.get("Content-Length", 0))
with open(dest, "wb") as of:
resp = RQ.get(url, stream=True)
for chunk in resp.iter_content(1024 * 1024):
of.write(chunk)
self.progress.emit(
{"done": of.tell(), "size": size, "outfile": outfile}
)
if not self.running:
return
class App(QApplication):
def __init__(self):
super().__init__(sys.argv)
self.setStyle("Fusion")
self.setup_styles()
def set_style(self, style):
print("LOAD:", style)
self.setPalette(self.styles[style])
def setup_styles(self):
self.styles = {}
styles = {
"Dark": {
"Window": QColor(53, 53, 53),
"WindowText": Qt.white,
"Base": QColor(15, 15, 15),
"AlternateBase": QColor(53, 53, 53),
"ToolTipBase": Qt.white,
"ToolTipText": Qt.white,
"Text": Qt.white,
"Button": QColor(53, 53, 53),
"ButtonText": Qt.white,
"BrightText": Qt.red,
"Highlight": QColor(255, 128, 0),
"HighlightedText": Qt.black,
}
}
for style, colors in styles.items():
palette = QPalette()
for entry, color in colors.items():
palette.setColor(getattr(QPalette, entry), color)
if color == Qt.darkGray:
palette.setColor(
QPalette.Disabled, getattr(QPalette, entry), QColor(53, 53, 53)
)
else:
palette.setColor(
QPalette.Disabled, getattr(QPalette, entry), Qt.darkGray
)
self.styles[style] = palette
self.styles["Light"] = self.style().standardPalette()
class WRoute(Ui_diag_route):
def __init__(self, main_window, hops):
super().__init__()
self.route = hops
dialog = QDialog(main_window)
self.setupUi(dialog)
for n, item in enumerate(hops):
if item["body"].startswith(item["system"]):
item["body"] = item["body"].replace(item["system"], "").strip()
item = QTreeWidgetItem(
self.lst_route,
[
str(n + 1),
item["system"],
"{body} ({star_type})".format(**item),
str(item["distance"]),
"<NotImplemented>", # Jump distance
],
)
item.setFlags(item.flags() & ~Qt.ItemIsDropEnabled)
dialog.exec_()
class ED_LRR(Ui_ED_LRR):
dl_thread = None
diag_prog = None
dl_started = None
system_found = pyqtSignal("PyQt_PyObject")
def __init__(self):
super().__init__()
self.current_job = None
def get_open_file(self, filetypes, callback=None):
fileName, _ = QFileDialog.getOpenFileName(
self.main_window,
"Open file",
cfg["folders.data_dir"],
filetypes,
options=QFileDialog.DontUseNativeDialog,
)
if callback:
return callback(fileName)
return fileName
def get_save_file(self, filetypes, callback=None):
fileName, _ = QFileDialog.getSaveFileName(
self.main_window,
"Save file",
cfg["folders.data_dir"],
filetypes,
options=QFileDialog.DontUseNativeDialog,
)
if callback:
return callback(fileName)
return fileName
def preprocess(self):
if self.current_job:
# ERROR
return
bodies_json = self.inp_bodies_pp.currentText()
systems_json = self.inp_systems_pp.currentText()
output_file = self.inp_out_pp.currentText()
self.current_job = PreprocessJob(
self.app, self.main_window, systems_json, bodies_json, output_file
)
def set_sys_lst(self, path):
if path not in cfg["history.stars_csv_path"]:
cfg["history.stars_csv_path"].append(path)
self.update_dropdowns()
def set_bodies_file(self, path):
if path not in cfg["history.bodies_path"]:
cfg["history.bodies_path"].append(path)
self.update_dropdowns()
def set_systems_file(self, path):
if path not in cfg["history.systems_path"]:
cfg["history.systems_path"].append(path)
self.update_dropdowns()
def update_dropdowns(self):
self.inp_systems_pp.clear()
self.inp_systems_dest_dl.clear()
for path in cfg["history.systems_path"][:]:
self.inp_systems_pp.addItem(path)
self.inp_systems_pp.setCurrentText(path)
self.inp_systems_dest_dl.addItem(path)
self.inp_systems_dest_dl.setCurrentText(path)
self.inp_bodies_pp.clear()
self.inp_bodies_dest_dl.clear()
for path in cfg["history.bodies_path"][:]:
self.inp_bodies_pp.addItem(path)
self.inp_bodies_pp.setCurrentText(path)
self.inp_bodies_dest_dl.addItem(path)
self.inp_bodies_dest_dl.setCurrentText(path)
self.inp_sys_lst.clear()
self.inp_out_pp.clear()
for path in cfg["history.stars_csv_path"]:
self.inp_sys_lst.addItem(path)
self.inp_sys_lst.setCurrentText(path)
self.inp_out_pp.addItem(path)
self.inp_out_pp.setCurrentText(path)
return
def log(self, *args):
t = datetime.today()
msg_t = "[{}] {}".format(t, str.format(*args))
self.txt_log.append(msg_t)
def set_comp_mode(self, _):
if self.rd_comp.isChecked():
comp_mode = "Compute Route"
self.btn_add.setText("Add")
if self.rd_precomp.isChecked():
comp_mode = "Precompute Graph"
self.btn_add.setText("Select")
self.log("COMP_MODE", comp_mode)
self.lst_sys.setEnabled(self.rd_comp.isChecked())
self.btn_rm.setEnabled(self.rd_comp.isChecked())
self.cmb_mode.setEnabled(self.rd_comp.isChecked())
self.chk_permute.setEnabled(self.rd_comp.isChecked())
self.lbl_keep.setEnabled(self.rd_comp.isChecked())
self.lbl_mode.setEnabled(self.rd_comp.isChecked())
self.chk_permute_keep_first.setEnabled(self.rd_comp.isChecked())
self.chk_permute_keep_last.setEnabled(self.rd_comp.isChecked())
self.set_route_mode(self.rd_precomp.isChecked() or None)
def set_route_mode(self, mode=None):
if mode == None:
mode = self.cmb_mode.currentText()
self.lbl_greedyness.setEnabled(mode == "A*-Search")
self.sld_greedyness.setEnabled(mode == "A*-Search")
def set_greedyness(self, value):
self.lbl_greedyness.setText("Greedyness Factor ({:.0%})".format(value / 100))
@property
def systems(self):
ret = []
for n in range(self.lst_sys.topLevelItemCount()):
ret.append(self.sys_to_dict(n))
return ret
def sys_to_dict(self, n):
header = [
self.lst_sys.headerItem().data(c, 0)
for c in range(self.lst_sys.headerItem().columnCount())
]
system = [
self.lst_sys.topLevelItem(n).data(c, 0)
for c in range(self.lst_sys.topLevelItem(n).columnCount())
]
ret = dict(zip(header, system))
ret["id"] = getattr(self.lst_sys.topLevelItem(n), "__id__", None)
ret.pop(None, None)
return ret
def error(self, msg):
QMessageBox.critical(self.main_window, "ED_LRR Error", msg)
def get_sys_list(self):
if not self.inp_sys_lst.currentText():
self.error("System list is required!")
return
path = pathlib.Path(self.inp_sys_lst.currentText())
if not path.exists():
self.error("System list does not exist, run download and preprocess first!")
return
return path
def compute_route(self):
self.bar_status.showMessage("Resolving systems...")
num_resolved = self.resolve_systems()
if num_resolved:
if (
QMessageBox.question(
self.main_window,
"ED_LRR",
"Resolved {} system(s), are the names correct?".format(
num_resolved
),
)
== QMessageBox.No
):
return
self.bar_status.clearMessage()
print(self.systems)
systems = [s["id"] for s in self.systems]
jump_range = self.sb_range.value()
if jump_range == 0:
self.error(
"Your jump range is set to zero, you're not going to get anywhere that way"
)
return
mode = self.cmb_mode.currentText()
primary = self.chk_primary.isChecked()
keep_first = self.chk_permute_keep_first.isChecked()
keep_last = self.chk_permute_keep_last.isChecked()
permute = self.chk_permute.isChecked()
greedyness = (
self.sld_greedyness.value() / 100
if self.sld_greedyness.isEnabled()
else None
)
path = self.get_sys_list()
if path is None:
return
precomp = None
path = str(path)
mode = {
"Breadth-First Search": "bfs",
"A*-Search": "astar",
"Greedy-Search": "greedy",
}[mode]
print(
systems,
jump_range,
None,
mode,
primary,
permute,
keep_first,
keep_last,
greedyness,
precomp,
path,
os.cpu_count() - 1,
)
if not self.current_job:
self.bar_status.showMessage("Computing Route...")
self.current_job = RouterJob(
self.app,
self.main_window,
systems,
jump_range,
None,
mode,
primary,
permute,
keep_first,
keep_last,
greedyness,
precomp,
path,
os.cpu_count() - 1,
)
else:
self.error("there is already a job running!")
def find_sys_by_names(self, names):
t_s = datetime.today()
if not self.get_sys_list():
return None
# TODO: start thread/subprocess
ret = _ed_lrr.find_sys(names, self.inp_sys_lst.currentText())
print("Took:", datetime.today() - t_s)
return ret
def resolve_systems(self):
# TODO: show spinner
names = []
nums = []
for n in range(self.lst_sys.topLevelItemCount()):
sys_id = getattr(self.lst_sys.topLevelItem(n), "__id__", None)
if sys_id is not None:
continue
names.append(self.sys_to_dict(n)["Name"])
nums.append(n)
if not names:
return 0
systems = self.find_sys_by_names(names)
if systems is None:
return
for i, name in zip(nums, names):
_, system = systems[name]
self.lst_sys.topLevelItem(i).setData(0, 0, system["system"])
self.lst_sys.topLevelItem(i).setData(1, 0, system["star_type"])
self.lst_sys.topLevelItem(i).__id__ = system["id"]
return len(names)
# diff, item = self.find_sys_by_name(name)
# print("Found", (diff, item))
def add_system(self):
name = self.inp_sys.text()
item = QTreeWidgetItem(self.lst_sys, [name, None])
item.resolved = False
item.setFlags(item.flags() & ~Qt.ItemIsDropEnabled)
def remove_system(self):
root = self.lst_sys.invisibleRootItem()
for item in self.lst_sys.selectedItems():
root.removeChild(item)
def dl_canceled(self):
if self.dl_thread:
print("Cancel!")
try:
self.dl_thread.progress.disconnect()
except TypeError:
pass
self.dl_thread.stop()
self.dl_thread.wait()
self.diag_prog.close()
self.dl_thread = None
self.diag_prog = None
self.dl_started = None
def handle_dl_progress(self, args):
filename = os.path.split(args["outfile"])[-1]
if self.diag_prog is None:
self.diag_prog = ProgressDialog("", "Cancel", 0, 1000, self.main_window)
if self.dl_thread:
self.diag_prog.canceled.connect(self.dl_canceled)
self.diag_prog.show()
t_elapsed = datetime.today() - self.dl_started
rate = args["done"] / t_elapsed.total_seconds()
remaining = (args["size"] - args["done"]) / rate
rate = round(rate, 2)
# print(rate, remaining)
try:
t_rem = timedelta(seconds=remaining)
except OverflowError:
t_rem = "-"
msg = "Downloading {} [{}/{}] ({}/s)\n[{}/{}]".format(
filename,
sizeof_fmt(args["done"]),
sizeof_fmt(args["size"]),
sizeof_fmt(rate),
t_round(t_elapsed),
t_round(t_rem),
)
self.diag_prog.setLabelText(msg)
self.diag_prog.setWindowTitle("Downloading EDSM Dumps")
self.diag_prog.setValue((args["done"] * 1000) // args["size"])
def run_download(self):
if self.dl_thread:
return
self.dl_started = datetime.today()
self.dl_thread = DownloadThread(
self.inp_systems_dl.currentText(),
self.inp_systems_dest_dl.currentText(),
self.inp_bodies_dl.currentText(),
self.inp_bodies_dest_dl.currentText(),
)
self.dl_thread.progress.connect(self.handle_dl_progress)
self.dl_thread.start()
print(".")
def update_permute_chk(self, state):
self.chk_permute_keep_first.setEnabled(state)
self.chk_permute_keep_last.setEnabled(state)
self.lbl_keep.setEnabled(state)
def setup_signals(self):
self.btn_download.clicked.connect(self.run_download)
self.set_greedyness(self.sld_greedyness.value())
self.cmb_mode.currentTextChanged.connect(self.set_route_mode)
self.rd_comp.toggled.connect(self.set_comp_mode)
self.rd_precomp.toggled.connect(self.set_comp_mode)
self.sld_greedyness.valueChanged.connect(self.set_greedyness)
self.btn_go.clicked.connect(self.compute_route)
self.btn_add.clicked.connect(self.add_system)
self.btn_rm.clicked.connect(self.remove_system)
self.btn_preprocess.clicked.connect(self.preprocess)
self.chk_permute.stateChanged.connect(self.update_permute_chk)
self.btn_out_browse_pp.clicked.connect(
lambda: self.get_save_file("CSV File (*.csv)", self.set_sys_lst)
)
self.btn_sys_lst_browse.clicked.connect(
lambda: self.get_open_file("CSV File (*.csv)", self.set_sys_lst)
)
self.btn_bodies_browse_pp.clicked.connect(
lambda: self.get_open_file("JSON File (*.json)", self.set_bodies_file)
)
self.btn_bodies_dest_browse_dl.clicked.connect(
lambda: self.get_save_file("JSON File (*.json)", self.set_bodies_file)
)
self.btn_systems_browse_pp.clicked.connect(
lambda: self.get_open_file("JSON File (*.json)", self.set_systems_file)
)
self.btn_systems_dest_browse_dl.clicked.connect(
lambda: self.get_save_file("JSON File (*.json)", self.set_systems_file)
)
def handle_close(self):
for key in [
"history.stars_csv_path",
"history.bodies_path",
"history.systems_path",
]:
for path in cfg[key][:]:
if not os.path.isfile(path):
cfg[key].remove(path)
cfg.write()
print("BYEEEEEE!")
def setup_styles(self, win, app):
for name in app.styles:
action = QAction(app)
action.setObjectName("action_load_style_" + name)
action.setText(name)
action.triggered.connect(lambda _, name=name: app.set_style(name))
self.menuStyle.addAction(action)
def setupUi(self, MainWindow, app):
super().setupUi(MainWindow)
self.update_dropdowns()
self.main_window = MainWindow
self.app = app
self.setup_signals()
self.lst_sys.setHeaderLabels(["Name", "Type"])
self.set_route_mode()
self.update_permute_chk(self.chk_permute.isChecked())
self.setup_styles(MainWindow, app)
def main():
MP.freeze_support()
app = App()
app.setWindowIcon(QIcon(r"D:\devel\rust\ed_lrr_gui\icon\icon.ico"))
MainWindow = QMainWindow()
MainWindow.setWindowIcon(QIcon(r"D:\devel\rust\ed_lrr_gui\icon\icon.ico"))
ui = ED_LRR()
ui.setupUi(MainWindow, app)
MainWindow.show()
ret = app.exec_()
ui.handle_close()
exit(ret)
if __name__ == "__main__":
main()