import hashlib
import json
from pathlib import Path

import numpy as np
from scipy.sparse import csr_matrix


class FixedNetwork:
    def __init__(self, layers, metadata):
        self.layers = layers
        self.metadata = metadata

    def forward(self, x, dense=False):
        x = np.asarray(x, dtype=np.float32)
        if x.ndim != 2 or x.shape[1] != 128:
            raise ValueError("Нужен массив [число позиций, 128]")
        for i, (w, b) in enumerate(self.layers):
            x = x @ w.toarray().T + b if dense else (w @ x.T).T + b
            if i + 1 < len(self.layers):
                x = np.maximum(x, 0)
        return x[:, 0]

    def save(self, path):
        path = Path(path)
        path.mkdir(parents=True, exist_ok=False)
        arrays = {}
        for i, (w, b) in enumerate(self.layers):
            for key, value in (("data", w.data), ("indices", w.indices), ("indptr", w.indptr), ("shape", w.shape), ("bias", b)):
                arrays[f"l{i}_{key}"] = np.asarray(value)
        np.savez(path / "weights.npz", **arrays)
        digest = hashlib.sha256((path / "weights.npz").read_bytes()).hexdigest()
        metadata = dict(self.metadata, sha256=digest, layers=len(self.layers), **self.stats())
        (path / "model.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n")

    @classmethod
    def load(cls, path):
        path = Path(path)
        meta = json.loads((path / "model.json").read_text())
        if hashlib.sha256((path / "weights.npz").read_bytes()).hexdigest() != meta["sha256"]:
            raise ValueError("Контрольная сумма весов не совпала")
        layers = []
        with np.load(path / "weights.npz", allow_pickle=False) as z:
            for i in range(meta["layers"]):
                w = csr_matrix((z[f"l{i}_data"], z[f"l{i}_indices"], z[f"l{i}_indptr"]), shape=tuple(z[f"l{i}_shape"]))
                layers.append((w, z[f"l{i}_bias"]))
        return cls(layers, meta)

    def stats(self):
        return {
            "architecture": [128] + [w.shape[0] for w, b in self.layers],
            "nonzero_weights": sum(w.nnz for w, b in self.layers),
            "nonzero_biases": sum(int(np.count_nonzero(b)) for w, b in self.layers),
            "dense_parameters_including_zeros": sum(w.shape[0] * w.shape[1] + len(b) for w, b in self.layers),
            "neurons_excluding_input": sum(w.shape[0] for w, b in self.layers),
            "sparse_array_bytes": sum(w.data.nbytes + w.indices.nbytes + w.indptr.nbytes + b.nbytes for w, b in self.layers),
            "float_weight_and_bias_bytes": sum(w.data.nbytes + b.nbytes for w, b in self.layers),
        }


class Layer:
    def __init__(self, inputs):
        self.inputs = inputs
        self.rows, self.cols, self.values, self.bias = [], [], [], []

    def node(self, terms, bias=0):
        row = len(self.bias)
        for col, value in terms:
            if value:
                self.rows.append(row)
                self.cols.append(col)
                self.values.append(value)
        self.bias.append(bias)
        return row

    def finish(self):
        w = csr_matrix((np.asarray(self.values, np.float32), (self.rows, self.cols)), shape=(len(self.bias), self.inputs), dtype=np.float32)
        w.sum_duplicates()
        w.eliminate_zeros()
        return w, np.asarray(self.bias, np.float32)