import re

from . import util


class Stats:
    def __init__(self):
        # Pepare some regex for parsing
        self.re_path_tracing = re.compile(
            r".*(Rendered|Path Tracing Tile) ([0-9]+)/([0-9]+)( Tiles)?" +
            r"(, Sample ([0-9]+)/([0-9]+))?$")
        self.re_total_render_time = re.compile(
            r".*Total render time: ([0-9]+(\.[0-9]+)?)")
        self.re_render_time_no_sync = re.compile(
            r".*Render time \(without synchronization\): ([0-9]+(\.[0-9]+)?)")
        self.re_pipeline_time = re.compile(
            r"Time: ([0-9:.]+) \(Saving: ([0-9:.]+)\)")
        self.re_cycles_memory = re.compile(
            r".*\| Mem:([0-9.]+[KM]?), Peak:([0-9.]+[KM]?) \|.*")

        # Render time stats.
        self.total_render_time = "N/A"
        self.render_time_no_sync = "N/A"
        self.pipeline_render_time = "N/A"

        # Render memory stats.
        self.device_peak_memory = "N/A"
        self.device_memory_usage = "N/A"

        # Current stats.
        self.path_tracing = ""
        self.current_tiles = 0
        self.total_tiles = 0
        self.current_sample = None
        self.total_samples = None

    def update(self, line):
        # Current tile progress.
        match = self.re_path_tracing.match(line)
        if match:
            self.path_tracing = match.group(1)
            self.current_tiles = int(match.group(2))
            self.total_tiles = int(match.group(3))
            self.current_sample = match.group(6)
            if self.current_sample is not None:
                self.current_sample = int(self.current_sample)
            self.total_samples = match.group(7)
            if self.total_samples is not None:
                self.total_samples = int(self.total_samples)
        # Total render time.
        match = self.re_total_render_time.match(line)
        if match:
            self.total_render_time = float(match.group(1))
        # Render time without sync.
        match = self.re_render_time_no_sync.match(line)
        if match:
            self.render_time_no_sync = float(match.group(1))
        # Total pipeline time.
        match = self.re_pipeline_time.match(line)
        if match:
            self.pipeline_render_time = \
                util.humanReadableTimeToSeconds(match.group(1))
        # Memory usage.
        match = self.re_cycles_memory.match(line)
        if match:
            mem = util.humanReadableSizeToMegabytes(match.group(1))
            peak = util.humanReadableSizeToMegabytes(match.group(2))
            if self.device_memory_usage == "N/A" or \
               mem > self.device_memory_usage:
                self.device_memory_usage = mem
            if self.device_peak_memory == "N/A" or \
               peak > self.device_peak_memory:
                self.device_peak_memory = peak

    def getCurrentProgress(self):
        if self.total_tiles == 0:
            return 0
        if not self.total_samples:
            return float(self.current_tiles) / float(self.total_tiles) * 100
        current_tile = self.current_tiles
        # Workaround for current master.
        if self.total_samples == self.current_sample:
            current_tile = max(current_tile - 1, 0)
        total_samples = self.total_tiles * self.total_samples
        current_sample = current_tile * self.total_samples + self.current_sample
        return float(current_sample) / float(total_samples) * 100

    def print(self):
        # TODO(sergey): Check that all stats are available.
        print("Total pipeline render time: {} ({} sec)"
              . format(util.humanReadableTimeDifference(
                  self.pipeline_render_time),
                  self.pipeline_render_time))
        print("Total Cycles render time: {} ({} sec)"
              . format(util.humanReadableTimeDifference(
                  self.total_render_time),
                  self.total_render_time))
        print("Pure Cycles render time (without sync): {} ({} sec)"
              . format(util.humanReadableTimeDifference(
                  self.render_time_no_sync),
                  self.render_time_no_sync))
        print("Cycles memoty usage: {} ({} peak)"
              . format(self.device_memory_usage,
                       self.device_peak_memory))

    def asDict(self):
        return {'total_render_time': self.total_render_time,
                'render_time_no_sync': self.render_time_no_sync,
                'pipeline_render_time': self.pipeline_render_time,
                'device_peak_memory': self.device_peak_memory,
                'device_memory_usage': self.device_memory_usage}