import init
import logging
from flask import Flask, request
import config
from config import app_host, app_port
from config import TIMEZONE, SPREADSHEETID, fb_PAGE_ACCESS_TOKEN, fb_VERIFY_TOKEN
from utils import createEvent, getGoogleSheetService, getEventsByPhone, getBookingDateFromEvent
from utils import getSheetValues,getClientFilter, getEventById, chunks, getGoogleCalendarService
from utils import getWeekDays, toTimestamp, toDatetime, updateSheet, utc2local, addMonths
from utils import listGet, p, userCacheGet, userCacheSet, setting, CLIENTS_CACHE_VALID
from datetime import datetime, timedelta
import time
from fbmq import Page, Template
import re
from googleapiclient.errors import HttpError
from lang import trans, strfweekday
import json

# fbmq page
page = Page(fb_PAGE_ACCESS_TOKEN)

# constant
VIEW_MY_BOOKING = 'VIEW_MY_BOOKING'
CANCEL_MY_BOOKING = 'CANCEL_MY_BOOKING'
CONFIRM_CANCEL_MY_BOOKING = 'CONFIRM_CANCEL_MY_BOOKING'
MAKE_A_BOOKING = 'MAKE_A_BOOKING'
CHOOSE_A_WEEK = 'CHOOSE_A_WEEK'
CHOOSE_A_DAY = 'CHOOSE_A_DAY'
CHOOSE_LANGUAGE = 'CHOOSE_LANGUAGE'

# init app
app = Flask(__name__)

def sendMsg(sender_id, name):
    return page.send(sender_id, trans(sender_id, name))

# actions
@app.route("/")
def hello():
    return "Hello World!"

@app.route("/create-events")
def createEvents():
    rows = getSheetValues()
    if rows:
        events = []
        for i, row in enumerate(rows[1:]):
            # Print columns A and E, which correspond to indices 0 and 4.
            # print('%s, %s' % (row[0], row[4]))
            if len(row) > 7 and row[4] and row[5] and row[6] and row[7]:
                if '/' in row[4]:
                    m, d, y = row[4].split('/')
                else:
                    y, m, d = row[4].split('-')
                y = int(y)
                m = int(m)
                d = int(d)
                time = datetime.now()
                time = time.replace(month=m, day=d, year=y, hour=int(
                    row[5]), minute=int(row[6]), second=0)
                events.append(['%s %s (%s)' %
                               (row[0], row[1], row[2]), time, float(row[7]), i])
                row[4] = ''
                row[5] = ''
                row[6] = ''
                row[7] = ''
        #
        for value in events:
            createEvent(*value[:-1])

        # save cleared
        if len(events) > 0:
            body = {
                'values': rows
            }
            updateSheet(body)
    return 'ok'

@app.route("/sheet-changed")
def sheetChanged():
    setting({CLIENTS_CACHE_VALID: False})
    return 'ok'

# facebook
@app.route('/webhook', methods=['GET'])
def validate():
    if request.args.get('hub.mode', '') == 'subscribe' and \
            request.args.get('hub.verify_token', '') == fb_VERIFY_TOKEN:

        print("Validating webhook")

        return request.args.get('hub.challenge', '')
    else:
        return 'Failed validation. Make sure the validation tokens match.'

@app.route('/webhook', methods=['POST'])
def webhook():
    requestStr = request.get_data().decode('utf-8')
    # e.g.: {"object":"page","entry":[{"id":"498812477183171","time":1522222835374,"messaging":[{"recipient":{"id":"498812477183171"},"timestamp":1522222835374,"sender":{"id":"1286804074753753"},"postback":{"payload":"CHOOSE_A_DAY_1522566413","title":"Sun (04-01)"}}]}]}
    # ignore encode error when facebook request data including utf8 chars
    try:
        page.handle_webhook(requestStr)
    except UnicodeEncodeError as e:
        pass
    except:
        sender_id = json.loads(requestStr).get('entry')[0]['messaging'][0]['sender']['id']
        sendMsg(sender_id, 'sys_error')
        logging.exception("Uncaught error in webhook")
    return "ok"

def sendButtons(to_id, title, buttons):
    title = trans(to_id, title)
    for chunked in chunks(buttons, 3):
        for btn in chunked:
            btn['title'] = trans(to_id, btn['title'])
        page.send(to_id, Template.Buttons(title, chunked))

def start(event):
    sender_id = event.sender_id
    locale = userCacheGet(sender_id, 'locale')
    if not locale:
        sendLanguagePicker(sender_id)
    else:
        buttons = [
            {
                "type": "postback",
                "value": VIEW_MY_BOOKING,
                "title": 'view_booking',
            },
            {
                "type": "postback",
                "value": CANCEL_MY_BOOKING,
                "title": 'cancel_booking',
            },
            {
                "type": "postback",
                "value": MAKE_A_BOOKING,
                "title": 'make_booking',
            },
        ]
        sendButtons(sender_id, 'hello', buttons)
def sendLanguagePicker(sender_id):
    buttons = [
            {
                "type": "postback",
                "value": CHOOSE_LANGUAGE + '_en',
                "title": "English"
            },
            {
                "type": "postback",
                "value": CHOOSE_LANGUAGE + '_zh',
                "title": "中文"
            },
        ]
    sendButtons(sender_id, 'pls_choose_language', buttons)
@page.callback([CHOOSE_LANGUAGE + '_(.+)'])
def choose_language(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    locale = payload[len(CHOOSE_LANGUAGE) + 1:]
    userCacheSet(sender_id, 'locale', locale)
    start(event)

@page.handle_message
def message_handler(event):
    """:type event: fbmq.Event"""
    sender_id = event.sender_id
    message = event.message_text
    print('receive msg', sender_id)
    if re.match(r'^\d\d?-\d\d?$', message):
        # date
        m,d = message.split('-')
        m = int(m)
        d = int(d)
        dt = utc2local(datetime.utcnow())
        try:
            dt = addMonths(dt, 1)
        except Exception as e:
            sendMsg(sender_id, 'invalid_date')
            return
        if dt.month != m:
            sendMsg(sender_id, 'date_not_next_month')
            return
        try:
            dt = dt.replace(day=d)
        except:
            sendMsg(sender_id, 'invalid_input')
            return
        setDate(dt, event)
    elif re.match(r'^\d\d?:\d\d?$', message):
        # time
        h,m = message.split(':')
        setTime([int(h), int(m)], event)
    else:
        t = message.lower()
        if 'language' in t or 'locale' in t or '语言' in t:
            sendLanguagePicker(sender_id)
        else:
            start(event)

@page.handle_referral
def handler2(event):
    """:type event: fbmq.Event"""
    sender_id = event.sender_id
    start(event)
    # page.send(sender_id, "thank you! your message is '%s'" % message)
    try:
        phone = event.referral['ref']
    except Exception as e:
        phone = None
    if phone:
        print('phone', phone)
        row = getClientFilter().filter(phone=phone).first()
        if row:
            if row.facebook_id:
                print('record already has facebook id.')
            else:
                print('store facebook id in google sheet')
                values = [
                    [
                        sender_id
                    ],
                    # Additional rows ...
                ]
                body = {
                    'values': values
                }
                rangeName = 'Sheet1!D%s:D%s' % (row.id, row.id)
                updateSheet(body, rangeName)
                print('facebook id stored in google sheet')
        else:
            print('no record found with given phone')
    print('no ref(phone) in hook event')

@page.callback([VIEW_MY_BOOKING])
def callback_1(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    row = getClientFilter().filter(facebook_id=sender_id).first()
    if row:
        events = getEventsByPhone(row.phone)
        if not events:
            sendMsg(sender_id, 'no_record')
            print('no booking record found', sender_id)
        else:
            bkdt = trans(sender_id, 'booking_date')
            ls = ['%s: %s'%(bkdt, getBookingDateFromEvent(event)) for event in events]
            page.send(sender_id, '\n'.join(ls))
    else:
        sendMsg(sender_id, 'no_record')

@page.callback([CANCEL_MY_BOOKING])
def callback_2(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    row = getClientFilter().filter(facebook_id=sender_id).first()
    if row:
        events = getEventsByPhone(row.phone)
        if not events:
            sendMsg(sender_id, 'no_record')
            print('no booking record found', sender_id)
        else:
            if len(events) == 1:
                bookingDatetime = getBookingDateFromEvent(events[0])
                bkdt = trans(sender_id, 'booking_date')
                bookingInfo = '%s: %s'%(bkdt, bookingDatetime)
                buttons = [
                    {
                        "type": "postback",
                        "value": CONFIRM_CANCEL_MY_BOOKING + '_' + str(event['id']),
                        "title": 'confirm_cancel'

                    },
                ]
                sendButtons(sender_id, bookingInfo, buttons)
            else:
                buttons = []
                for event in events:
                    bookingDatetime = getBookingDateFromEvent(event, '%m-%d %H:%M')
                    bookingInfo = trans(sender_id, 'booking') + ': ' + bookingDatetime
                    buttons.append({
                        "type": "postback",
                        "value": CONFIRM_CANCEL_MY_BOOKING + '_' + str(event['id']),
                        "title": bookingInfo,
                    })

                sendButtons(sender_id, 'choose_cancel', buttons)
    else:
        sendMsg(sender_id, 'no_record')
        print('no booking record found', sender_id)

@page.callback([CONFIRM_CANCEL_MY_BOOKING + '_(.+)'])
def callback_3(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    evid = payload[len(CONFIRM_CANCEL_MY_BOOKING) + 1:]
    event = getEventById(evid)
    if not event:
        sendMsg(sender_id, 'no_record')
        print('no booking record found', sender_id)
    else:
        bookingDatetime = getBookingDateFromEvent(event)
        service = getGoogleCalendarService()
        try:
            service.events().delete(calendarId='primary', eventId=event['id']).execute()
        except HttpError as e:
            if e.resp.status in [410]:
                # already deleted
                pass
            else:
                raise e
        time.sleep(1)
        event = getEventById(evid)
        if event:
            sendMsg(sender_id, 'failed_cancel')
            print('Failed to cancel the booking', sender_id, event)
        else:
            page.send(sender_id, trans(sender_id, 'cancelled_successfully')%(bookingDatetime))
            print('Booking canceled successfully', sender_id, event)

@page.callback([MAKE_A_BOOKING])
def callback_4(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    buttons = []
    titles = ['this_week', 'next_week', 'week_after_next', 'week_after_next2', 'next_month']
    for i in range(5):
        buttons.append({
            "type": "postback",
            "value": CHOOSE_A_WEEK + '_' + str(i),
            "title": titles[i]
        })
    sendButtons(sender_id, 'pls_choose', buttons)

@page.callback([CHOOSE_A_WEEK + r'_(\d)'])
def callback_5(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    n = int(payload[len(CHOOSE_A_WEEK) + 1:])
    if n < 4:
        # week
        t = datetime.utcnow() + timedelta(weeks=n)
        weekdays = getWeekDays(t)
        buttons = [{"type": "postback", "value": CHOOSE_A_DAY + '_' + str(toTimestamp(v)), "title": strfweekday(sender_id, v, '%a (%m-%d)')} for v in weekdays]
        sendButtons(sender_id, 'pls_choose', buttons)
    else:
        # next month
        sendMsg(sender_id, 'pls_input_date')
@page.callback([CHOOSE_A_DAY + r'_(\d+)'])
def callback_6(payload, event):
    sender_id = event.sender_id
    print(payload, sender_id)
    n = int(payload[len(CHOOSE_A_DAY) + 1:])
    dt = toDatetime(n)
    return setDate(dt, event)

def setDate(dt, event):
    sender_id = event.sender_id
    print('set date', sender_id)
    userCacheSet(sender_id, 'date', str(toTimestamp(dt)))
    sendMsg(sender_id, 'pls_input_time')

def setTime(timeList, event):
    sender_id = event.sender_id
    print('set time', sender_id)
    t = userCacheGet(sender_id, 'date')
    if not t:
        sendMsg(sender_id, 'pls_set_date_before_time')
        return
    dt = toDatetime(int(t))
    h, m = timeList
    try:
        dt = dt.replace(hour=h, minute=m, second=0, microsecond=0)
    except:
        sendMsg(sender_id, 'invalid_input')
        return
    row = getClientFilter().filter(facebook_id=sender_id).first()
    if not row:
        sendMsg(sender_id, 'no_record')
        return
    createEvent('%s %s (%s)' %(row.phone, row.name, row.full_name), dt)
    sendMsg(sender_id, 'booked_successfully')

@page.after_send
def after_send(payload, response):
    """:type payload: fbmq.Payload"""
    print("complete")

# bootstrap app
if __name__ == '__main__':
    app.run(host=app_host, port=app_port, debug=True, threaded=True)


# Mateusz Łuczyński 331826
def d(s: str) -> int:
    n = len(s)
    return n * n

def alg(s: str, dict, maxS) -> int:
    n = len(s)
    dp = [(0, None)] + [(-1, None) for x in range(n-1)]
    for x in range(n):
        y = max(x - maxS + 1, 0)
        while x - y + 1 > 0:
            sub = s[y : (x+1)]
            if sub in dict:
                weight = d(sub)
                if y == 0:
                    pref = 0
                    val = weight
                else:
                    pref = dp[y-1][0]
                    val = dp[y-1][0] + weight
                if pref != -1 and val > dp[x][0]:
                        dp[x] = (val, y-1)
            y += 1
    res = ""
    backtrack = dp[n-1][1]
    end = n
    while True:
        if backtrack == None:
            return None
        sub = s[(backtrack+1) : end]
        res = sub + " " + res
        if backtrack == -1:
            break
        end = backtrack+1
        backtrack = dp[backtrack][1]
    return res

def main():
    f = open("polish_words.txt", "r", encoding="utf-8")
    maxS = -1
    dict = {}
    for word in f:
        word = word.strip()
        n = len(word)
        if n > maxS:
            maxS = n
        dict[word] = True
    f.close()
    inp = open("zad2_input.txt", "r", encoding="utf-8")
    out = open("zad2_output.txt", "w", encoding="utf-8")
    for line in inp:
        line = line.strip()
        s = alg(line, maxS, dict)
        out.write(s + '\n')
    inp.close()
    out.close()

if __name__ == "__main__":
    main()


from selenium import webdriver

# Create ChromeOptions and add arguments as needed
chrome_options = webdriver.ChromeOptions()
chrome_options.add_argument('--headless')  # Run in headless mode to avoid opening a visible browser window

# Initialize the Chrome WebDriver with the specified options
driver = webdriver.Chrome(options=chrome_options)

# Get the path to the ChromeDriver executable
driver_path = driver.service.process.pid

# Print the path to ChromeDriver
print("Path to ChromeDriver executable:", driver_path)

# Quit the WebDriver
driver.quit()


import random
import struct
from cStringIO import StringIO

import pylzma

from edacc.constants import JOB_STATUS_COLOR, JOB_RESULT_CODE_COLOR

def newline_split_string(s, n):
    """ Splits the string s up into parts of size n and returns
        the concatenation of these parts separated by newline characters
        as result. """
    if n == 0: return s
    return '\n'.join([s[i:i + n] for i in range(0, len(s), n)])

def download_size(value):
    """ Takes an integer number of bytes and returns a pretty string representation """
    if value <= 0:
        return "0 Bytes"
    elif value < 1024:
        return str(value) + " Bytes"
    elif value < 1024 * 1024:
        return "%.1f kB" % (value / 1024.0)
    else:
        return "%.1f MB" % (value / 1024.0 / 1024.0)

def job_status_color(value):
    """ Returns an HTML conform color string for the job status """
    if value not in JOB_STATUS_COLOR:
        return 'grey'
    else:
        return JOB_STATUS_COLOR[value]

def job_result_code_color(value):
    """ Returns an HTML conform color string for the job result code """
    if str(value).startswith("1"):
        return "#00CC33"
    if value not in JOB_RESULT_CODE_COLOR:
        return 'grey'
    else:
        return JOB_RESULT_CODE_COLOR[value]

def parameter_string(solver_config):
    """ Returns a string of the solver configuration parameters """
    parameters = solver_config.parameter_instances
    args = []
    for p in parameters:
        args.append(p.parameter.prefix or "")
        if p.parameter.name == 'seed':
            args.append("<seed>")
        if p.parameter.name == 'instance':
            args.append("<instance>")
        if p.parameter.name == 'tempdir':
            args.append("<tempdir>")
        if p.parameter.hasValue:
            if p.value == "": # if value not set, use default value from parameters table
                args.append(p.parameter.defaultValue or "")
            else:
                args.append(p.value or "")
    return " ".join(args)

def parameter_template(solver):
    """ Returns a string of the solver configuration parameters """
    parameters = solver.parameters
    args = []
    for p in sorted(parameters, key=lambda p: p.order):
        if p.prefix: args.append(p.prefix)
        if p.space and p.prefix and p.hasValue: args.append(" ")
        if p.name == 'seed':
            args.append("SEED")
        if p.name == 'tempdir':
            args.append("TEMPDIR")
        if p.name == 'instance':
            args.append("INSTANCE")
        if p.hasValue:
            args.append(p.defaultValue or "")
        args.append(" ")
    return "".join(args)

def result_time(time):
    """ Returns a representation of the time value. If time is None
        a dash character is returned, otherwise the time is returned """
    if time is None:
        return '-'
    else:
        return time

def launch_command(solver_config):
    """ Returns a string of what the solver launch command looks like
    given the solver configuration
    """
    return (solver_config.solver_binary.runCommand or "") + " ./" + (
    solver_config.solver_binary.runPath or "") + " " + parameter_string(solver_config)

def datetimeformat(value, format='%H:%M, %d-%m-%Y'):
    """ Returns the passed datetime value as formatted string according to the formatting string
    :format:
    """
    if value is None: return ""
    return value.strftime(format)

def competition_phase(value):
    """ Returns a textual label of a competiton phase given by an integer value """
    if value == 1:
        return "Category Definition Phase"
    elif value == 2:
        return "Registration and Submission Phase"
    elif value == 3:
        return "Solver Testing Phase"
    elif value == 4:
        return "Solver Resubmission Phase"
    elif value == 5:
        return "Competition Phase"
    elif value == 6:
        return "Release Phase"
    elif value == 7:
        return "Post-Release Phase"
    else:
        return "unknown phase"

def parse_parameters(parameters):
    """ Parse parameters from the solver submission form, returns a list
        of tuples (name, prefix, default_value, boolean, order)
    """
    parameters = parameters.strip().split()
    params = []
    i = 0
    while i < len(parameters):
        if parameters[i].startswith('-') and not "=" in parameters[i]:
            # prefixed parameter
            if i + 1 < len(parameters) and (parameters[i + 1].lower() in ('seed', 'instance', 'tempdir')):
                pname = parameters[i + 1].lower()
                prefix = parameters[i]
                default_value = ''
                boolean = False
                params.append((pname, prefix, default_value, boolean, i, True))
                i += 2
            else:
                pname = parameters[i]
                prefix = parameters[i]
                if i + 1 == len(parameters) or parameters[i + 1].startswith('-'):
                    boolean = True
                    default_value = ''
                    params.append((pname, prefix, default_value, boolean, i, True))
                    i += 1
                else:
                    boolean = False
                    default_value = parameters[i + 1]
                    params.append((pname, prefix, default_value, boolean, i, True))
                    i += 2
        else:
            if "=" in parameters[i] and (
                        "INSTANCE" in parameters[i] or "SEED" in parameters[i] or "TEMPDIR" in parameters[i]):
                pname = parameters[i].split("=")[1].lower()
                prefix = parameters[i].split("=")[0] + "="
                default_value = ''
                boolean = False
                space = False
            else:
                # parameter without prefix
                if parameters[i].lower() in ('seed', 'instance', 'tempdir'):
                    pname = parameters[i].lower()
                    prefix = ''
                    default_value = ''
                    boolean = False
                    space = True
                else:
                    pname = parameters[i]
                    prefix = parameters[i]
                    default_value = ''
                    boolean = True
                    space = True
            params.append((pname, prefix, default_value, boolean, i, space))
            i += 1
    return params

# some SAT functions
random.seed()

def random_clause(l):
    return [random.randint(0, 1) for _ in xrange(l)]

def random_formula(clauses, clauseLength):
    return [random_clause(clauseLength) for _ in xrange(clauses)]

def assignment(n):
    if n == 1:
        yield [0]
        yield [1]
    else:
        for a in assignment(n - 1):
            yield a + [0]
            yield a + [1]

def satisfies(a, f):
    sat_clauses = 0
    for clause in f:
        for i in xrange(len(clause)):
            if clause[i] == a[i]: sat_clauses += 1; break
    if sat_clauses == len(f):
        return True
    return False

def SAT(f):
    for a in assignment(len(f[0])):
        if satisfies(a, f): return a
    return None

def render_formula(f):
    res = []
    for c in f:
        cl = []
        for i in xrange(len(c)):
            if c[i] == 0:
                cl.append(u'\u00ac' + chr(i + ord('A')))
            else:
                cl.append(chr(i + ord('A')))
        res.append('(' + u' \u2228 '.join(cl) + ')')
    return u' \u2227 '.join(res)

def formatOutputFile(data):
    if data is not None:
        if len(data) > 4 * 1024:
            # show only the first and last 2048 characters if the resultFile is larger than 4kB
            resultFile_text = data[:2048] + "\n\n... [truncated " + str(
                int((len(data) - 4096) / 1024.0)) + " kB]\n\n" + data[-2048:]
        else:
            resultFile_text = data
    else:
        resultFile_text = "No output"
    return resultFile_text

def lzma_decompress(data):
    """
        LZMA decompression using pylzma.
        The LZMA header consists of 5 + 8 bytes.
        The first 5 bytes are compression parameters, the 8 following bytes specify
        the length of the data.
    """
    # unpack the data length from the LZMA header (bytes # 6-13 inclusively/unsigned long long int)
    coded_length = struct.unpack('<Q', data[5:13])[0]
    if coded_length == '\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF':
        # if the length is -1 (in two's complement), there is an EOS marker
        # marking the end of the data and pylzma doesn't need the coded_length
        return pylzma.decompress(data[:5] + data[13:])
    else:
        # if the length is specified, pylzma needs the coded_length since there probably
        # is no EOS marker
        return pylzma.decompress(data[:5] + data[13:], maxlength=coded_length)

def lzma_compress(data):
    """ LZMA compression using pylzma """
    c = pylzma.compressfile(StringIO(data), dictionary=8, fastBytes=128,
                            algorithm=0, eos=0, matchfinder='hc3')
    result = c.read(5)
    result += struct.pack('<Q', len(data))
    return result + c.read()

def truncate_name(s, l=100):
    if len(s) > l:
        return s[:l / 2] + " [..] " + s[-l / 2:]
    return s


#-------------------------------------------------------------------------------
# Program: blink.py
# Author: schnelled (Dustin Schnelle)
# Date: 2/22/19
# Description: Turns on an LED for one second, then off for one second, repeat
#-------------------------------------------------------------------------------
from nanpy import (ArduinoApi, SerialManager)
from time import sleep

# Setup the variable (LED)
LED = 13

# Attempt to connect the Arduino
try:
    # Connect and create the Arduino
    connection = SerialManager()
    a = ArduinoApi(connection = connection)
# Otherwise it failed to connect
except:
    # Print error message
    print("Error: Failed to connect to Arduino!")

# Setup the pinMode as if we were in the Arduino IDE
a.pinMode(LED, a.OUTPUT)

try:
    # Continuous loop
    while True:
        # Turn on the LED
        a.digitalWrite(LED, a.HIGH)
        # Wait one second
        sleep(1)
        # Turn off the LED
        a.digitalWrite(LED, a.LOW)
        # Wait one second
        sleep(1)
except:
    # Turn off the LED
    a.digitalWrite(LED, a.LOW)


# 5 Autos
# 2 Garagen
# Garagen müssen wissen welche Autos sie haben

auto1 = ["Audi", 11]
auto2 = ["BMW", 12]
auto3 = ["Mercedes", 13]
auto4 = ["Lambo", 14]
auto5 = ["Ferrari", 25000]

garage_1 = [auto1, auto2, auto3]
garage_2 = [auto4, auto5]

# Function übernimmt den Paramater Garage und printet alle Autos

def print_garage(some_garage):
    for auto in some_garage:
        print(f'{auto[0]} hat {auto[1]} Kilometer')

for auto in garage_1:
    auto[1] += 10000

print_garage(garage_1)
print_garage(garage_2)


from django.test import TestCase
from rest_framework import status
from rest_framework.reverse import reverse
from rest_framework.test import APITestCase
from rest_framework.utils import json

from managebook.models import Book
from managebook.serialize import BookSerializer

class TestBook(APITestCase):
    def test_create(self):
        url = reverse("create")
        data = {
            "title": "title text",
            "text": "text ttt"
        }
        response = self.client.post(url, data)
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)
        count_book = Book.objects.all().count()
        self.assertEqual(count_book, 1)

    def test_delete(self):
        book = Book.objects.create(title="test t", text="tes;odjf")
        url = reverse("delete", kwargs={"id": book.id})
        response = self.client.delete(url)
        self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
        count_book = Book.objects.all().count()
        self.assertEqual(count_book, 0)

    def test_list(self):
        book1 = Book(title="test t", text="tes;odjf")
        book2 = Book(title="test t2", text="tes;odjf2")
        Book.objects.bulk_create([book1, book2])
        url = reverse("list")
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.json()[0]['title'], "test t")
        self.assertEqual(response.json()[1]['title'], "test t2")

    # def test_update(self):
    #     book = Book.objects.create(title="test t", text="tes;odjf")
    #     data = {
    #         "title": "title text",
    #         "text": "text ttt"
    #     }
    #     url = reverse("update", kwargs={id: book.id})
    #     response = self.client.put(url, json.dumps(data))
    #     self.assertEqual(response.status_code, )

    def test_update(self):
        book = Book.objects.create(title="test t", text="tes;odjf")
        data = {
            "title": "title text",
            "text": "text ttt"
        }
        url = reverse("update", kwargs={"id": book.id})
        response = self.client.put(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)


from abc import ABCMeta, abstractmethod
from math import fsum
import numpy as np
from scipy.stats import norm as norm_dist
from collections import namedtuple
from pyx import _hmmc
from hmm.multivariatenormal import MultivariateNormal, MixtureModel

__author__ = 'eranroz'

class HMMModel(object):
    """
    base model for HMM
    """
    __metaclass__ = ABCMeta

    def __init__(self, state_transition, emission, min_alpha=None):
        """
        Initializes a new HMM model.
        @param state_transition: state transition matrix.
                    (A & Pi in Rabiner's paper)
                    with rows - source state, cols - target state.
                    0 state assumed to be the begin state
                    (according to Durbin's book)
        @param emission: observation symbol probability distribution
                        (B in Rabiner's paper)
                        rows - states, cols - output
                        The begin state should have emission too (doesn't matter what)

        @param min_alpha: prior on the sequence length (transition to itself)
        """

        self.state_transition = state_transition
        self.emission = emission
        self.min_alpha = min_alpha

    def num_states(self):
        """
        Get number of states in the model
        """
        return self.state_transition.shape[0]

    def num_alphabet(self):
        """
        Get number of symbols in the alphabet
        """
        return self.emission.shape[1]

    def get_state_transition(self):
        """
        State transition matrix: rows - source state, cols - target state
        """
        return self.state_transition

    def get_emission(self):
        """
        Emission matrix:  rows states, cols - symbols
        """
        return self.emission

    @abstractmethod
    def _maximize_emission(self, seq, gammas):
        """
        part of the maximization step of the EM algorithm (Baum-Welsh)
        should update the emission probabilities according to the forward-backward results

        @param seq symbol sequence
        @param gammas from backward forward
        """
        pass

    def collect_stats(self, seq, bf_output):
        """
        Collect statistics based on the expectations, for later use in the maximization step
        @param seq: observed sequence
        @param bf_output: backward forward output
        @return: statistics on transitions, statistics on emission
        """
        transition_stats = self._collect_transition_stats(seq, bf_output)
        emission_stats = self._collect_emission_stats(seq, bf_output.state_p)
        return transition_stats, emission_stats

    def maximize_using_stats(self, transition_stats, emission_stats):
        """
        Use collected statistics to maximize the model.
        @param transition_stats: statistics of the transitions
        @param emission_stats: statistics on emission
        """
        self._maximize_transition_stats(transition_stats)
        self._maximize_emission_stats(emission_stats)

    @abstractmethod
    def _collect_emission_stats(self, seq, gammas):
        """
        collects statistics from backward forward iteration (without normalization) about emission
        @param seq: observation sequence
        @param gammas: matrix of states to probabilities
        @return:
        """
        pass

    def _maximize_transition_stats(self, transition_stats, prob_w):
        # == maximize transition ==

        prob_w = np.ones(prob_w.shape[0])  # replace prob_w with ones - all are equal (it is not normalized)
        transition_stats = np.average(transition_stats, 0, prob_w)
        transition_stats[1:, 1:] /= np.sum(transition_stats[1:, 1:], 1)[:, None]  # normalize
        if self.min_alpha is not None:
            n_states = transition_stats.shape[0] - 1  # minus begin state
            diagonal_selector = np.eye(n_states, dtype='bool')
            self_transitions = transition_stats[1:, 1:][diagonal_selector]
            n_self_transitions = np.maximum(self.min_alpha, self_transitions)
            # reduce the diff from the rest of transitions equally
            transition_stats[1:, 1:][~diagonal_selector] -= (n_self_transitions - self_transitions) / (n_states - 1)
            transition_stats[1:, 1:][diagonal_selector] = n_self_transitions

        # start transition
        transition_stats[0, 1:] /= np.sum(transition_stats[0, 1:])

        # end transition
        transition_stats[1:, 0] /= np.sum(transition_stats[1:, 0])
        # update transition matrix
        self.state_transition = transition_stats

    @abstractmethod
    def _maximize_emission_stats(self, emission_stats, prob_w):
        """
        collects statistics from backward forward iteration (without normalization) about emission
        @param emission_stats: statistics for emission of each state
        @return:
        """
        pass

    def _collect_transition_stats(self, seq, bf_output):
        """
        collects statistics from backward forward iteration (without normalization) about transitions
        @param seq: observation sequence
        @param bf_output: output of backward forward iteration
        @return:
        """

        # collect statistics for the transition matrix
        new_state_transition = self.state_transition.copy()
        emission = self.get_emission()
        back_emission_seq = emission[1:, seq].T
        back_emission_seq *= bf_output.backward / bf_output.scales[:, None]

        new_state_transition[1:, 1:] *= np.dot(bf_output.forward[:-1, :].transpose(), back_emission_seq[1:, :])
        # we avoid normalization in the collection phase:
        # new_state_transition[1:, 1:] /= np.sum(new_state_transition[1:, 1:], 1)[:, None]

        # start transition
        # new_state_transition[0, 1:] = bf_output.state_p[0, :]
        new_state_transition[0, 1:] *= back_emission_seq[0, :]
        new_state_transition[0, 1:] /= np.sum(new_state_transition[0, 1:])

        # end transition
        new_state_transition[1:, 0] = bf_output.forward[-1, :] * back_emission_seq[-1, :] / bf_output.scales[-1]
        # we normalize it although it is collect stats because there must be one
        new_state_transition[1:, 0] /= np.sum(new_state_transition[1:, 0])
        return new_state_transition

    def _maximize_transition(self, seq, bf_output):
        """
        part of the maximization step of the EM algorithm (Baum-Welsh)
        should state transition probabilities according to the forward-backward results

        @param seq: observation sequence
        @param bf_output: output of the forward-backward algorithm
        @return:
        """
        new_state_transition = self.state_transition.copy()
        emission = self.get_emission()
        back_emission_seq = emission[1:, seq].T
        back_emission_seq *= bf_output.backward / bf_output.scales[:, None]

        new_state_transition[1:, 1:] *= np.dot(bf_output.forward[:-1, :].transpose(), back_emission_seq[1:, :])
        new_state_transition[1:, 1:] /= np.sum(new_state_transition[1:, 1:], 1)[:, None]  # normalize
        if self.min_alpha is not None:
            n_states = new_state_transition.shape[0] - 1  # minus begin state
            diagonal_selector = np.eye(n_states, dtype='bool')
            self_transitions = new_state_transition[1:, 1:][diagonal_selector]
            n_self_transitions = np.maximum(self.min_alpha, self_transitions)
            # reduce the diff from the rest of transitions equally
            new_state_transition[1:, 1:][~diagonal_selector] -= (n_self_transitions - self_transitions) / (n_states - 1)
            new_state_transition[1:, 1:][diagonal_selector] = n_self_transitions

        # start transition
        new_state_transition[0, 1:] = bf_output.state_p[0, :]
        # new_state_transition[0, 1:] *= back_emission_seq[0, 1:]
        new_state_transition[0, 1:] /= np.sum(new_state_transition[0, 1:])

        # end transition
        new_state_transition[1:, 0] = bf_output.forward[-1, :] * back_emission_seq[-1, :] / bf_output.scales[-1]
        new_state_transition[1:, 0] /= np.sum(new_state_transition[1:, 0])
        # update transition matrix
        self.state_transition = new_state_transition

    def maximize(self, seq, bw_output):
        """
        Maximization step for in Baum-Welsh algorithm (EM)

        @param seq symbol sequence
        @param bw_output results of backward forward (scaling version)
        """
        self._maximize_transition(seq, bw_output)
        self._maximize_emission(seq, bw_output.state_p)

    def __str__(self):
        return '\n'.join(
            ['Model parameters:', 'Emission:', str(self.emission), 'State transition:',
             str(self.state_transition)])

    def viterbi(self, symbol_seq):
        """
        Find the most probable path through the model

        Dynamic programming algorithm for decoding the states.
        Implementation according to Durbin, Biological sequence analysis [p. 57]

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        """

        n_states = self.num_states()
        unique_values = set(symbol_seq)
        emission_seq = np.zeros((len(symbol_seq), n_states - 1))
        for v in unique_values:
            emission_seq[symbol_seq == v, :] = np.log(self.get_emission()[1:, v])

        return _hmmc.viterbi(emission_seq, self.state_transition)

    def viterbi_old(self, symbol_seq):
        """
        Find the most probable path through the model

        Same as above, but not optimized. just for simplification or if you got to trouble with compilation

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        """
        n_states = self.num_states()
        unique_values = set(symbol_seq)
        emission_seq = np.zeros((len(symbol_seq), n_states - 1))
        for v in unique_values:
            emission_seq[symbol_seq == v, :] = np.log(self.get_emission()[1:, v])

        ptr_mat = np.zeros((len(symbol_seq), n_states - 1))
        l_state_trans_mat_T = self.get_state_transition().T

        emission_iterator = iter(emission_seq)
        ptr_iterator = iter(ptr_mat)
        # initial condition is begin state
        prev = next(emission_iterator) + 1 + np.log(l_state_trans_mat_T[1:, 0])
        next(ptr_iterator)[...] = np.argmax(prev)
        end_state = 0  # termination step
        end_transition = np.log(l_state_trans_mat_T[end_state, 1:])
        l_state_trans_mat_T = np.log(l_state_trans_mat_T[1:, 1:])

        # recursion step
        for emission_symbol in emission_iterator:
            p_state_transition = prev + l_state_trans_mat_T
            max_k = np.max(p_state_transition, 1)
            next(ptr_iterator)[...] = np.argmax(p_state_transition, 1)
            prev = emission_symbol + max_k

        p_state_transition = prev + end_transition

        # traceback step and without begin state
        most_probable_path = np.zeros(len(symbol_seq), int)
        most_probable_path[-1] = np.argmax(p_state_transition)

        for i in np.arange(len(symbol_seq) - 1, 0, -1):
            most_probable_path[i - 1] = ptr_mat[i, most_probable_path[i]]

        return most_probable_path

    def forward_backward(self, symbol_seq, model_end_state=False, num_stable=False):
        """
        Calculates the probability for the model and each step in it

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        @param model_end_state: whether to consider end state or not
        @param num_stable: whether to handle numerical stability by changing inf to max number (may cause slowness)
        Remarks:
        this implementation uses scaling variant to overcome floating points errors.
        """
        n_states = self.num_states()
        emission = self.get_emission()

        emission_seq = emission[1:, symbol_seq].T
        state_trans_mat = self.get_state_transition()
        dot = np.dot  # shortcut for performance
        real_transitions_T = state_trans_mat[1:, 1:].T.copy(order='C')
        real_transitions_T2 = state_trans_mat[1:, 1:].copy(order='C')

        # emission * transition
        e_trans = iter(emission_seq[1:, :, np.newaxis] * real_transitions_T)
        forward_iterator = np.nditer([emission_seq, np.newaxis, np.newaxis],
                                     flags=['external_loop', 'reduce_ok'],
                                     op_flags=[
                                         ['readonly'],
                                         ['readwrite', 'allocate', 'no_broadcast'],
                                         ['readwrite', 'allocate', 'no_broadcast']
                                         ],
                                     op_axes=[[-1, 0, 1], [-1, 0, 1], [-1, 0, -1]], order='C')

        # -----	  forward algorithm	  -----
        # initial condition is begin state (in Durbin there is another forward - the begin = 1)
        tup = next(forward_iterator)  # emission_i, forward_i,scaling_i
        try:
            tup[1][...] = np.maximum(state_trans_mat[0, 1:] * tup[0], 1e-100)
        except FloatingPointError:
            tup[1][...] = np.exp(np.maximum(np.log(state_trans_mat[0, 1:])+np.log(tup[0]), -100))
        tup[2][...] = fsum(tup[1])
        tup[1][...] /= tup[2]
        prev_forward = tup[1]

        # recursion step
        for tup in forward_iterator:
            prev_forward = dot(next(e_trans), prev_forward)

            # scaling - see Rabiner p. 16, or Durbin p. 79
            scaling = tup[2]
            scaling[...] = fsum(prev_forward)  # fsum is more numerical stable
            tup[1][...] = prev_forward = prev_forward / scaling

        forward = forward_iterator.operands[1]
        s_j = forward_iterator.operands[2]
        # end transition
        log_p_model = np.sum(np.log(s_j))
        if model_end_state:  # Durbin - with end state
            end_state = 0  # termination step
            end_transition = forward[emission_seq.shape[0] - 1, :] * state_trans_mat[1:, end_state]
            log_p_model += np.log(sum(end_transition))

        # backward algorithm
        # initial condition is end state
        if model_end_state:
            prev_back = (state_trans_mat[1:, 0])  # Durbin p.60
        else:
            prev_back = np.ones(n_states - 1)  # Rabiner p.7 (24)

        backward_iterator = np.nditer([emission_seq[:0:-1], np.newaxis],  # / s_j[:, None]
                                      flags=['external_loop', 'reduce_ok'],
                                      op_flags=[['readonly'],
                                                ['readwrite', 'allocate']],
                                      op_axes=[[-1, 0, 1], [-1, 0, 1]], order='C')

        # recursion step
        e_trans = iter((emission_seq / s_j[:, np.newaxis])[:0:-1, np.newaxis, :] * real_transitions_T2)\

        # keep the scale if we get out of numerical boundaries (may occur if the transition matrix has 0)
        if num_stable:
            for tup in backward_iterator:  # emission_i/scale_i, backward_i
                dot(next(e_trans), prev_back, tup[1])  # = dot(real_transitions_T2, prev_back * tup[0])
                prev_back = tup[1][...] = np.nan_to_num(prev_back)
        else:
            for tup in backward_iterator:  # emission_i/scale_i, backward_i
                dot(next(e_trans), prev_back, tup[1])  # = dot(real_transitions_T2, prev_back * tup[0])

        if model_end_state:
            backward = np.append(backward_iterator.operands[1][::-1], state_trans_mat[1:, 0][np.newaxis, :],
                                 axis=0)  # Durbin p.60
        else:
            backward = np.append(backward_iterator.operands[1][::-1], np.ones((1, n_states - 1)),
                                 axis=0)  # Rabiner p.7 (24)

        # return bf_result
        bf_result = namedtuple('BFResult', 'model_p state_p forward backward scales')
        return bf_result(log_p_model, backward * forward, forward, backward, s_j)

    def forward_backward_old(self, symbol_seq, model_end_state=False):
        """
        Calculates the probability for the model and each step in it

        Same as above but with less optimizations. for easier reading
        @param model_end_state: whether to model the end state
        @param symbol_seq: sequence of symbols/observations

        """
        n_states = self.num_states()
        emission = self.get_emission()
        state_trans_mat = self.get_state_transition()

        s_j = np.ones(len(symbol_seq))
        forward = np.zeros((len(symbol_seq), n_states - 1), order='F')  # minus the begin state
        backward = np.zeros((len(symbol_seq), n_states - 1), order='F')

        # forward algorithm
        # initial condition is begin state (in Durbin there is another forward - the begin = 1)
        forward[0, :] = state_trans_mat[0, 1:] * emission[1:, symbol_seq[0]]
        s_j[0] = sum(forward[0, :])
        forward[0, :] /= s_j[0]
        prev_forward = forward[0, :]
        # recursion step
        # transform to emission array
        unique_values = set(symbol_seq)
        emission_seq = np.zeros((len(symbol_seq), n_states - 1))
        for v in unique_values:
            emission_seq[symbol_seq == v, :] = emission[1:, v]

        real_transitions = state_trans_mat[1:, 1:]
        t_real_transitions = real_transitions.transpose()
        p_state_transition = np.zeros(n_states - 1)
        summing_arr = np.array([1] * (n_states - 1))
        emission_iterator = iter(emission_seq)
        next(emission_iterator)  # skip the first (instead of condition in for loop
        s_j_iterator = np.nditer(s_j, op_flags=['writeonly'])
        forward_iterator = np.nditer(forward, flags=['external_loop'], op_flags=['writeonly'], order='C')
        next(s_j_iterator)
        next(forward_iterator)
        for sym_emission in emission_iterator:
            np.dot(t_real_transitions, prev_forward, p_state_transition)
            prev_forward = sym_emission * p_state_transition

            # scaling - see Rabiner p. 16, or Durbin p. 79
            scaling = np.dot(summing_arr, prev_forward)  # dot is actually faster then np.sum(prev_forward)
            next(s_j_iterator)[...] = scaling
            next(forward_iterator)[...] = prev_forward = prev_forward / scaling

        # end transition
        log_p_model = np.sum(np.log(s_j))
        if model_end_state:  # Durbin - with end state
            end_state = 0  # termination step
            end_transition = forward[len(symbol_seq) - 1, :] * state_trans_mat[1:, end_state]
            log_p_model += np.log(sum(end_transition))

        # backward algorithm
        # initial condition is end state
        if model_end_state:
            prev_back = backward[len(symbol_seq) - 1, :] = (state_trans_mat[1:, 0])  # Durbin p.60
        else:
            prev_back = backward[len(symbol_seq) - 1, :] = [1, 1]  # Rabiner p.7 (24)

        backward_iterator = np.nditer(backward[::-1], flags=['external_loop'], op_flags=['writeonly'], order='C')
        next(backward_iterator)
        s_j_iterator = iter(s_j[::-1])

        # recursion step
        for sym_emission in emission_seq[:0:-1]:
            np.dot(prev_back * sym_emission, t_real_transitions, p_state_transition)
            next(backward_iterator)[...] = prev_back = p_state_transition / next(s_j_iterator)

        bf_result = namedtuple('BFResult', 'model_p state_p forward backward scales')
        return bf_result(log_p_model, backward * forward, forward, backward, s_j)

    def forward_backward_log(self, symbol_seq, model_end_state=False):
        """
        Calculates the probability for the model and each step in it

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        @param model: an HMM model to calculate on the given symbol sequence
        @param model_end_state: whether to consider end state or not

        Remarks:
        this implementation uses log variant to overcome floating points errors instead of scaling.
        """
        interpolation_res = 0.01
        interpolation_res_i = 1.0 / interpolation_res
        interpol_tbl = np.log(1 + np.exp(-np.arange(0, 35, interpolation_res)))
        last_interpol = len(interpol_tbl) - 1

        def interpolate(prev):
            """
            Uses interpolation table to calculate log r=log(p+log(1+exp(q-p))) which is approx equals to
            log r = log (max)+(exp(1+log(x)) [x=min-max from prev].
            see Durbin p. 79
            @param prev: previous probabilities plus the transition
            @return: result of interpolation of log r=log(p+log(1+exp(q-p)))
            """
            maxes = np.max(prev, 1)
            interpolation_i = np.minimum(np.round(-interpolation_res_i * (np.sum(prev, 1) - 2 * maxes)),
                                         last_interpol).astype(int)
            return maxes + interpol_tbl[interpolation_i]

        l_emission = np.log(self.get_emission()[1:, :])
        forward = np.zeros((len(symbol_seq), self.num_states() - 1))
        backward = np.zeros((len(symbol_seq), self.num_states() - 1))

        l_state_transition = np.log(self.get_state_transition()[1:, 1:])
        l_t_state_transition = l_state_transition.transpose()
        # forward algorithm
        # initial condition is begin state (in Durbin there is another forward - the begin = 1)
        prev_forward = forward[0, :] = np.log(self.get_state_transition()[0, 1:]) + l_emission[:, symbol_seq[0]]
        # recursion step
        emission_seq = list(enumerate([l_emission[:, s] for s in symbol_seq]))

        for i, sym_emission in emission_seq:
            if i == 0:
                continue

            from_prev = prev_forward + l_t_state_transition  # each row is different state, each col - CHECKED
            # now the sum approximation
            from_prev = interpolate(from_prev)
            forward[i, :] = prev_forward = sym_emission + from_prev

        # termination step
        if model_end_state:
            end_state = 0
            log_p_model = forward[len(symbol_seq) - 1, :] + np.log(self.get_state_transition()[1:, end_state])
            log_p_model = interpolate(np.array([log_p_model]))
        else:
            log_p_model = interpolate([forward[len(symbol_seq) - 1, :]])

        # backward algorithm
        last_index = len(symbol_seq) - 1
        if model_end_state:
            prev_back = backward[last_index, :] = np.log(self.get_state_transition()[1:, 0])
        else:
            prev_back = backward[last_index, :] = [0, 0]  # Rabiner p.7 (24)

        for i, sym_emission in reversed(emission_seq):
            if i == 0:
                continue
            p_state_transition = interpolate(l_state_transition + (prev_back + sym_emission))
            prev_back = backward[i - 1, :] = p_state_transition

        # posterior probability 3.14 (P.60 Durbin)
        l_posterior = forward + backward - log_p_model
        bf_result = namedtuple('BFResult', 'model_p state_p forward backward')
        return bf_result(log_p_model, l_posterior, forward, backward)

    def html_state_transition(self):
        """
        nice html representation (as table) of state transition matrix
        """
        import matplotlib as mpl
        import matplotlib.cm as cm
        backgrounds = cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=np.min(self.state_transition), vmax=np.max(self.state_transition)), cmap=cm.Blues)
        color_mapper = lambda x:  'rgb(%i, %i, %i)' % backgrounds.to_rgba(x, bytes=True)[:3]

        cells_ths = ''.join(['<th>%i</th>' % i for i in np.arange(1, self.state_transition.shape[0])])
        states_trs = []
        for state_i, state_trans in enumerate(self.state_transition[1:, 1:]):
            state_description = "<td style=\"font-weight:bold;\">%i</td>" % (state_i+1)
            state_description += ''.join(['<td style="color:#fff;background:%s">%.2f</td>' % (color_mapper(val), val)
                                          for val in state_trans])

            states_trs.append('<tr>%s</tr>' % state_description)

        template = """
<table style="font-size:85%;text-align:center;border-collapse:collapse;border:1px solid #aaa;" cellpadding="5" border="1">
<tr style="font-size:larger; font-weight: bold;">
    <th>State/Cell</th>
    {cells_ths}
</tr>
{states_trs}
</table>
"""
        return template.format(**({'cells_ths': cells_ths,
                                   'states_trs': '\n'.join(states_trs)}))

class DiscreteHMM(HMMModel):
    """
    Discrete Hidden Markov Model

     Handles sequences with discrete alphabet
    """
    def _collect_emission_stats(self, seq, gammas):
        new_emission_matrix = np.zeros((self.num_states()-1, self.num_alphabet()))  # except begin state

        state_p = gammas
        for sym in range(0, self.num_alphabet()):
            where_sym = (seq == sym)
            new_emission_matrix[:, sym] = np.sum(state_p[where_sym, :], 0)

        return new_emission_matrix

    def _maximize_emission_stats(self, emission_stats, prob_w):
        # normalize
        prob_w = np.ones(prob_w.shape[0])  # replace to 1 - all equals (non normalized)
        emission_stats = np.average(emission_stats, 0, prob_w)
        emission_stats /= np.sum(emission_stats, 1)[:, None]
        self.emission[1:, :] = emission_stats

    def _maximize_emission(self, seq, gammas):
        new_emission_matrix = self._collect_emission_stats(seq, gammas)

        # normalize
        new_emission_matrix /= np.sum(new_emission_matrix, 1)[:, None]

        self.emission[1:, :] = new_emission_matrix

class ContinuousHMM(HMMModel):
    """
    Continuous HMM for observations of real values

    The states are gaussian (or gaussian mixtures)
    @param state_transition: state transition matrix
    @param mean_vars: array of mean, var tuple (or array of such for mixtures)
    @param emission_density: log-concave or elliptically symmetric density
    @param mixture_coef: mixture coefficients
    """

    def __init__(self, state_transition, mean_vars, emission_density=norm_dist, mixture_coef=None,
                 min_alpha=None):
        emission = _ContinuousEmission(mean_vars, emission_density, mixture_coef)
        super().__init__(state_transition, emission, min_alpha=min_alpha)

    def _maximize_emission(self, seq, gammas):
        mean_vars = np.zeros((self.num_states(), 2))
        min_std = 0.5
        if self.emission.mixtures is None:
            state_norm = np.sum(gammas, 0)
            mu = np.sum(gammas * seq[:, None], 0) / state_norm
            old_mu = self.emission.mean_vars[1:, 0]
            sym_min_mu = np.power(seq[:, None] - old_mu, 2)
            std = np.sqrt(np.sum(gammas * sym_min_mu, 0) / state_norm)

            std = np.maximum(std, min_std)  # it must not get to zero
            mean_vars[1:, :] = np.column_stack([mu, std])
            self.emission = _ContinuousEmission(mean_vars, self.emission.dist_func)
        else:  # TODO: not yet fully tested
            mean_vars = [(0, 0)]
            mixture_coeff = [1]
            for state in np.arange(0, self.num_states() - 1):
                has_coeff = True
                try:
                    if len(self.emission.mixtures[state + 1]) > 1:
                        coeff_pdfs = [self.emission.dist_func(mean, var).pdf for mean, var in
                                      self.emission.mean_vars[state + 1]]
                        coeff_obs = np.array([[p(s) for p in coeff_pdfs] for s in seq])
                        coeff_obs /= np.sum(coeff_obs, 1)[:, None]
                        gamma_coeff = coeff_obs * gammas[:, state][:, None]
                        seq_man = seq[:, None]
                    else:
                        gamma_coeff = gammas[:, state]
                except TypeError:
                    gamma_coeff = gammas[:, state]
                    seq_man = seq
                    has_coeff = False

                sum_gamma = np.sum(gamma_coeff, 0)

                mu = np.sum(gamma_coeff * seq_man, 0) / sum_gamma
                mu *= self.emission.mixtures[state + 1]
                sym_min_mu = np.power(seq_man - mu, 2)
                std = np.sqrt(np.sum(gamma_coeff * sym_min_mu, 0) / sum_gamma)
                std = np.maximum(std, min_std)  # it must not get to zero
                if has_coeff:
                    mean_vars.append(list(zip(mu, std)))
                else:
                    mean_vars.append((mu, std))
                mixture_coeff.append(sum_gamma / np.sum(sum_gamma))

            self.emission = _ContinuousEmission(mean_vars, self.emission.dist_func, mixture_coeff)

class _ContinuousEmission():
    """
    Emission for continuous HMM.
    """

    def __init__(self, mean_vars, dist=norm_dist, mixture_coef=None):
        """
        Initializes a new continuous distribution states.
        @param mean_vars: np array of mean and variance for each state
        @param dist: distribution function
        @return: a new instance of ContinuousDistStates
        """
        self.dist_func = dist
        self.mean_vars = mean_vars
        self.mixtures = mixture_coef
        self.cache = dict()

        self.states = self._set_states()
        self._set_states()

    def _set_states(self):
        from functools import partial

        if self.mixtures is None:
            states = ([self.dist_func(mean, var).pdf for mean, var in self.mean_vars])
        else:
            states = []
            for mean_var, mixture in zip(self.mean_vars, self.mixtures):
                try:
                    mix_pdf = [self.dist_func(mean, var).pdf for mean, var in mean_var]
                    mix = partial(_ContinuousEmission.mixture_pdf, mix_pdf, mixture)

                    if np.abs(np.sum(mixture) - 1) > 1e-6:
                        raise Exception("Bad mixture - mixture for be summed to 1")
                except TypeError:
                    mix = self.dist_func(mean_var[0], mean_var[1]).pdf
                states.append(mix)
        return states

    @staticmethod
    def mixture_pdf(pdfs, mixtures, val):
        """
        Mixture distrbution
        @param pdfs:
        @param mixtures:
        @param val:
        @return:
        """
        return np.dot([p(val) for p in pdfs], mixtures)

    def __getitem__(self, x):
        """
        Get emission for state
        @param x:  first index is state (or slice for all states), second is value or array of values
        @return: p according to pdf
        """
        min_p = np.finfo(float).eps  # 1e-100
        if isinstance(x[0], slice):
            if isinstance(x[1], np.ndarray):  # this new case improves performance if you give emission array of values
                pdfs = np.array([dist(x[1]) for dist in self.states[x[0]]])
                pdfs = np.maximum(pdfs, min_p)
                return pdfs
            else:
                try:
                    return self.cache[x[1]]
                except KeyError:
                    pdfs = np.array([dist(x[1]) for dist in self.states[x[0]]])
                    pdfs = np.maximum(pdfs, min_p)
                    self.cache[x[1]] = pdfs
                    return self.cache[x[1]]
        else:
            return self.states[x[0]].pdf(x[1])

    def __getstate__(self):
        return {
            'mean_vars': self.mean_vars,
            'mixture_coef': self.mixtures,
            'dist_func': self.dist_func
        }

    def __setstate__(self, state):
        self.mean_vars = state['mean_vars']
        self.mixtures = state['mixture_coef']
        self.dist_func = state['dist_func']
        self.cache = dict()
        self.states = self._set_states()

    def __str__(self):
        return '\n'.join([str(self.dist_func.name) + ' distribution', 'Mean\t Var', str(self.mean_vars)])

class GaussianHMM(HMMModel):
    """
    Gaussian HMM for multidimensional gaussian mixtures. Extension for Continuous HMM above

    The states are gaussian mixtures
    @param state_transition: state transition matrix
    @param mean_vars: array of mean, var tuple (or array of such for mixtures)
    @param mixture_coef: mixture coefficients

    """
    def __init__(self, state_transition, mean_vars, mixture_coef=None, min_alpha=None):
        if len(mean_vars) == len(state_transition):
            mean_vars = mean_vars[1:]  # trim the begin emission

        # if not mixture (only tuple) wrap it with another list
        mean_vars = [[mean_cov] if isinstance(mean_cov, tuple) else mean_cov for mean_cov in mean_vars]

        # same type: all mean vars should be lists
        emission = _GaussianEmission(mean_vars, mixture_coef)
        super().__init__(state_transition, emission, min_alpha=min_alpha)

    def _collect_emission_stats(self, seq, gammas):
        state_norm = np.sum(gammas, 0)
        mean_vars = []
        mixture_coeff = []
        for state in np.arange(0, self.num_states() - 1):
            is_mixture = len(self.emission.mixtures[state]) > 1
            if is_mixture:
                emissions = self.emission.components_emission(state, seq)
                sum_emissions = np.sum(emissions, 0)
                emissions /= sum_emissions[:, None]  # normalize
                gamma_state = emissions * gammas[:, state][:, None]
                del emissions
                mixture_coeff.append(sum_emissions/np.sum(sum_emissions))
            else:
                gamma_state = gammas[:, state][:, None]
                mixture_coeff.append([1])

            covars_mixture = []
            new_means = np.dot(seq, gamma_state).T  # avoid normalization by state_norm[state]).T
            for mixture_i, mixture in enumerate(self.emission.mixtures[state]):
                gamma_c = gamma_state[:, mixture_i]
                old_mean = self.emission.mean_vars[state][mixture_i][0]
                seq_min_mean = seq - old_mean.T
                new_cov = np.dot((seq_min_mean * gamma_c), seq_min_mean.T)  # avoid normalization by state_norm[state])
                covars_mixture.append(new_cov)
            mean_vars.append((new_means, covars_mixture))

        return mean_vars, mixture_coeff, state_norm

    def _maximize_emission_stats(self, emissions_stats, prob_w):
        prob_w = np.ones_like(prob_w)  # we use non normalized terms(normalizinig by state_norms)
        # extract means and covariances
        mean_vars, mixture_coeff, state_norms = zip(*emissions_stats)
        state_norms = np.sum(state_norms, 0)
        # extract mixture coefficents
        mixture_coeff = np.array(mixture_coeff)
        new_mixcoeff = []
        min_std = 0.5  # np.finfo(float).eps  #1e-5 #
        new_mean_vars = []
        for state in np.arange(0, self.num_states() - 1):
            new_mixcoeff.append(np.average(mixture_coeff[:, state], 0, prob_w))

            mix_mean_covar = [mean_var_i[state] for mean_var_i in mean_vars]
            mean_state_i, covar_state_i = zip(*mix_mean_covar)

            mean_state_i = np.average(mean_state_i, 0, prob_w)[0]/state_norms[state]
            covar_state_i = np.average(covar_state_i, 0, prob_w)[0]/state_norms[state]

            # the diagonal must be large enough
            np.fill_diagonal(covar_state_i, np.maximum(covar_state_i.diagonal().copy(), min_std))
            new_mean_vars.append([(mean_state_i, covar_state_i)])

        self.emission = _GaussianEmission(new_mean_vars, new_mixcoeff)

    def _maximize_emission(self, seq, gammas):
        min_std = 0.5  # np.finfo(float).eps  #1e-5 #
        state_norm = np.sum(gammas, 0)
        mean_vars = []
        mixture_coeff = []
        for state in np.arange(0, self.num_states() - 1):
            is_mixture = len(self.emission.mixtures[state]) > 1
            if is_mixture:
                emissions = self.emission.components_emission(state, seq)
                sum_emissions = np.sum(emissions, 0)
                emissions /= sum_emissions[:, None]  # normalize
                gamma_state = emissions * gammas[:, state][:, None]
                del emissions
                mixture_coeff.append(sum_emissions/np.sum(sum_emissions))
            else:
                gamma_state = gammas[:, state][:, None]
                mixture_coeff.append([1])

            covars_mixture = []
            new_means = (np.dot(seq, gamma_state) / state_norm[state]).T
            for mixture_i, mixture in enumerate(self.emission.mixtures[state]):
                gamma_c = gamma_state[:, mixture_i]
                old_mean = self.emission.mean_vars[state][mixture_i][0]

                seq_min_mean = seq - old_mean.T
                new_cov = np.dot((seq_min_mean * gamma_c), seq_min_mean.T) / state_norm[state]
                # the diagonal must be large enough
                np.fill_diagonal(new_cov, np.maximum(new_cov.diagonal().copy(), min_std))
                covars_mixture.append(new_cov)
            mean_vars.append(list(zip(new_means, covars_mixture)))

        self.emission = _GaussianEmission(mean_vars, mixture_coeff)

    def viterbi(self, symbol_seq):
        """
        Find the most probable path through the model

        Dynamic programming algorithm for decoding the states.
        Implementation according to Durbin, Biological sequence analysis [p. 57]

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        """
        emission_seq = np.log(self.get_emission()[1:, symbol_seq]).T
        return _hmmc.viterbi(emission_seq, self.state_transition)

    def __str__(self):
        # handelding mixtures isn't handled currently
        means = np.array([x[0][0] for x in self.emission.mean_vars])
        covars = np.array([x[0][1].diagonal() for x in self.emission.mean_vars])

        str_rep = 'GMM. Means:\n'
        str_rep += np.array_str(means, precision=2, suppress_small=True, max_line_width=250).replace('\n\n', '\n')
        str_rep += '\nDiagonals for covariance matrices:\n'
        str_rep += np.array_str(covars, precision=2, suppress_small=True, max_line_width=250).replace('\n\n', '\n')
        str_rep += '\nStates transitions% (begin not shown):\n'
        str_rep += np.array_str(100*(self.state_transition[1:, 1:]), precision=1, suppress_small=True,
                                max_line_width=250)
        return str_rep

class _GaussianEmission():
    """
    Emission for continuous HMM.
    """

    def __init__(self, mean_vars, mixture_coef=None):
        """
        Initializes a new continuous distribution states.
        @param mean_vars: np array of mean and variance for each state
        @return: a new instance of ContinuousDistStates
        """
        # if no mixture defined set to 1
        if mixture_coef is None:
            mixture_coef = np.ones((len(mean_vars), 1))
        self.mean_vars = _GaussianEmission._normalize_mean_vars(mean_vars)
        self.mixtures = mixture_coef
        self.states, self.pseudo_states = self._set_states()

    @staticmethod
    def _normalize_mean_vars(mean_vars):
        # use numpy arrays
        n_mean_vars = []
        for state in mean_vars:
            state_mean_cov = []
            for mean, cov in state:
                mean = np.array(mean)
                if mean.ndim == 0:
                    mean = mean[None, None]
                if mean.ndim == 1:
                    mean = mean[None]
                cov = np.array(cov)
                if cov.ndim == 0:
                    cov = cov[None, None]
                if mean.ndim == 1:
                    cov = cov[None]
                state_mean_cov.append((mean, cov))
            n_mean_vars.append(state_mean_cov)
        return n_mean_vars

    def _set_states(self):
        states = []
        pseudo_states = []
        for mean_var, mixture in zip(self.mean_vars, self.mixtures):
            if mixture == 1:
                emission = MultivariateNormal(mean_var[0][0], mean_var[0][1])
                pseudo_states.append(emission.pdf)
            else:
                emission = MixtureModel([MultivariateNormal(mean, cov) for mean, cov in mean_var], mixture)
                pseudo_states.append(emission.components_pdf)
            states.append(emission.pdf)

        return states, pseudo_states

    def components_emission(self, states, observations, use_log=False):
        min_p = np.finfo(float).eps if use_log else -np.inf
        if isinstance(states, int):
            return np.maximum(self.pseudo_states[states].log_pdf(observations), min_p, use_log=use_log)
        pdfs = np.array([dist(observations, use_log=True) for dist in self.pseudo_states[states]]).T
        return np.maximum(pdfs, min_p)

    def __getitem__(self, x):
        """
        Get emission for state
        @param x:  first index is state (or slice for all states), second is value or array of values
        @return: p according to pdf
        """
        min_p = np.finfo(float).eps
        if isinstance(x[0], int):
            return np.maximum(self.states[x[0]].pdf(x[1]), min_p)

        pdfs = np.array([dist(x[1]) for dist in self.states])

        return np.maximum(pdfs, min_p)

    def __getstate__(self):
        return {
            'mean_vars': self.mean_vars,
            'mixture_coef': self.mixtures
        }

    def __setstate__(self, state):
        self.mean_vars = state['mean_vars']
        self.mixtures = state['mixture_coef']
        self.states, self.pseudo_states = self._set_states()

    def __str__(self):
        return 'Mean\t Var\n %s' % str(self.mean_vars)

class MultinomialHMM(HMMModel):
    """
    A multinomial HMM supporting multiple features
    @param state_transition: transition matrix between states
    @param emission_matrix: n X m X k matrix, n - state, m - feature , k - emission for char k
    @param min_alpha:
    """
    def __init__(self, state_transition, emission_matrix, min_alpha=None):
        emission = MultinomialEmission(emission_matrix)
        self.num_features = emission_matrix.shape[1]
        super().__init__(state_transition, emission, min_alpha=min_alpha)

    def _maximize_emission(self, seq, gammas):
        """
        num_features = seq.shape[1]
        alphabet_size = self.emission.states_prob.shape[2]
        new_emission_matrix = np.zeros((self.num_states(), num_features, alphabet_size))

        state_p = gammas
        for feature in range(seq.shape[1]):
            for sym in range(0, alphabet_size):
                where_sym = (seq[:, feature] == sym)
                new_emission_matrix[1:, feature, sym] = np.sum(state_p[where_sym, :], 0)

        # normalize
        new_emission_matrix[1:, :, :] /= np.sum(new_emission_matrix[1:, :, :], -1)[:, :, np.newaxis]

        self.emission = MultinomialEmission(new_emission_matrix)
        """
        emission_stats = self._collect_emission_stats(seq, gammas)
        self._maximize_emission_stats(emission_stats)

    def _collect_emission_stats(self, seq, gammas):
        num_features = seq.shape[1]
        alphabet_size = self.emission.states_prob.shape[2]
        new_emission_matrix = np.zeros((self.num_states(), num_features, alphabet_size))

        state_p = gammas
        for feature in range(seq.shape[1]):
            for sym in range(0, alphabet_size):
                where_sym = (seq[:, feature] == sym)
                new_emission_matrix[1:, feature, sym] = np.sum(state_p[where_sym, :], 0)

        return new_emission_matrix

    def _maximize_emission_stats(self, emission_stats, prob_w):
        # normalize
        prob_w = np.ones(prob_w.shape[0])  # replace to 1 - all equals (non normalized)
        emission_stats = np.average(emission_stats, 0, prob_w)
        emission_stats[1:, :, :] /= np.sum(emission_stats[1:, :, :], -1)[:, :, np.newaxis]
        self.emission = MultinomialEmission(emission_stats)

    def num_alphabet(self):
        """
        Number of symbols in alphabet based on teh defined emission (in all the features)

        @return:
        """
        return self.emission.states_prob.shape[1]*self.emission.states_prob.shape[2]

    def viterbi(self, symbol_seq):
        """
        Find the most probable path through the model

        Dynamic programming algorithm for decoding the states.
        Implementation according to Durbin, Biological sequence analysis [p. 57]

        @param symbol_seq: observed sequence (array). Should be numerical (same size as defined in model)
        """
        emission_seq = np.log(self.get_emission()[1:, symbol_seq]).T
        return _hmmc.viterbi(emission_seq, self.state_transition)

class MultinomialEmission():
    """
    MultinomialEmission assumes multinomial distribution e.g (n!/(n1!n2!))*(p1)^n1*(p2)^n2...
    """
    def __init__(self, states_prob):
        """
        initializes a new instance of MultinomialEmission
        @param states_prob: a matrix of nxmxk where n - number of states and m-feature index, k -emission for char k
        """
        self.states_prob = states_prob

    def __getitem__(self, item):
        """
        Get emission for state
        @param item:  first index is state (or slice for all states),
                      second is value or matrix of values [observations x features]
        @return: p according to pdf
        """
        min_p = np.finfo(float).eps
        if isinstance(item[0], int):
            features_prob = self.states_prob[item[0], :, item[1]]
            p = np.prod(features_prob, -1)
        else:

            # features_prob = self.states_prob[:, :, item[1]]
            # p = np.prod(features_prob, -1)
            features_prob = []
            for f in range(self.states_prob.shape[1]):
                features_prob.append(self.states_prob[item[0], f, item[1][:, f]])
            p = np.prod(features_prob, 0)
        p = np.maximum(p, min_p)
        return p


import os, sys
import pywavefront
import pickle

def upDir(x):
   return os.path.dirname(x)

def getMVSDir():
    return os.path.dirname(os.path.realpath(__file__))

sys.path.append(upDir(upDir(getMVSDir())))
from dnnlib.util import EasyDict

from cameras.camera import *
from graphics_utils.image_io import *
from graphics_utils.render_utils import transformMesh
from cameras.camera import  getAnisotropicScalingMatrix
from multiview.COLMAP.colmap_parser import parseColmapData
import glob
#=======================================================

# import cameras from bundler format
# https://www.cs.cornell.edu/~snavely/bundler/bundler-v0.4-manual.html
def camerasFromBundler(bundlerFile, resolution, modelMatrix=None):
    file = open(bundlerFile, 'r')
    cameras = []
    file.readline()
    camCount = int(file.readline().split()[0])

    # handle potentially different resolutions for MVS input views
    if not isinstance(resolution, list):
        resolution = [resolution] * camCount

    for i in range(camCount):
        camera = Camera()
        focalLength = float(file.readline().split()[0])
        rot = np.eye(3)
        for j in range(3):
            rot[j] = [ float(l) for l in file.readline().split() ]
        transl =  [ float(l) for l in file.readline().split() ]
        viewMatrix = np.eye(4)
        viewMatrix[0:3, 0:3] = rot
        viewMatrix[0:3, 3] = transl
        camera.setViewMatrix(viewMatrix)
        if modelMatrix is not None:
            camera = transformCamera(camera, modelMatrix)
        res = resolution[i]
        aspect = res[0] / res[1]
        fov = focalLengthToFov(focalLength, res[1])
        camera.setProjectionParams(fov, aspect)
        cameras.append(camera)
    file.close()
    return cameras

#=======================================================

# export cameras to bundler format
# https://www.cs.cornell.edu/~snavely/bundler/bundler-v0.4-manual.html
def camerasToBundler(outputFile, cameras, resolution, modelMatrix=None):
    file = open(outputFile, "w")
    file.write("# Bundle file v0.3\n")
    file.write(str(len(cameras)) + " 0\n")
    for cam in cameras:
        focalLength = fovToFocalLength(cam.fov, resolution[1])
        txt = str(focalLength) + " 0 0\n"
        if modelMatrix is not None:
            cam = transformCamera(cam, modelMatrix)
        v = cam.viewMatrix
        for r in range(3):
            for c in range(3):
                txt += str(v[r, c]) + " "
            txt += "\n"
        for c in range(3):
            txt += str(v[c, 3]) + " "
        txt += "\n"
        file.write(txt)
    file.close()

#=======================================================

# convert cameras to lookat format
def camerasToLookat(outputFile, cameras):
    file = open(outputFile, "w")
    for idx, cam in enumerate(cameras):
        txt = "Cam" + str(idx).zfill(3) + " -D origin="
        for p in range(3):
            txt += str(cam.position[p])
            if p < 2:
                txt += ","
        v = viewDirectionFromViewMatrix(cam.viewMatrix)
        target = cam.position + v
        txt += " -D target="
        for p in range(3):
            txt += str(target[p])
            if p < 2:
                txt += ","
        up = upVectorFromViewMatrix(cam.viewMatrix)
        txt += " -D up="
        for p in range(3):
            txt += str(up[p])
            if p < 2:
                txt += ","
        txt += " -D fovy="
        txt += str(np.degrees(cam.fov))
        txt += " -D clip=1.0,1000.0"
        txt += "\n"
        file.write(txt)
    file.close()

#=======================================================

# functions for importing data from Colmap or RealityCapture

# import scene from RealityCapture
# https://www.capturingreality.com/
def importRealityCaptureData(dataPath, useBlurred):

    #--------------------------------------------------------------
    # turn wavefront obj into mesh
    def expandWavefrontGeometry(geometry):
        vertices = np.array(list(geometry.materials.items())[0][1].vertices)  #vertex_format is "T2F_V3F"
        vertices = vertices.reshape((int(vertices.shape[0] / 5), 5))
        vertices = np.concatenate((vertices[..., 2:], vertices[..., :2]), axis=1)  #desired format is XYZUV
        mesh = EasyDict()
        mesh.vertices = vertices[..., :3].flatten()
        mesh.uvs = vertices[..., 3:].flatten()
        mesh.verticesAndUVs = vertices.flatten()
        mesh.faces = np.arange(int(vertices.shape[0])).flatten()
        return mesh
    #--------------------------------------------------------------

    modelPath = os.path.join(dataPath, "realitycapture", "model", "")
    registrationPath = os.path.join(dataPath, "realitycapture", "registration", "")

    # load mesh
    geometry = loadMesh(modelPath)
    mesh = expandWavefrontGeometry(geometry)

    # load texture (set it as Red if not found any)
    mesh.texture = loadTexture(modelPath, "RealityCapture")

    # load undistorted images
    images = loadUndistortedImages(registrationPath, useBlurred)

    # load camera
    cameraSetups = loadCameraSetup(registrationPath, images)

    # ignore blacklist images
    images, cameraSetups = excludeBlacklist(registrationPath, images, cameraSetups)

    return mesh, cameraSetups, images

# import scene from Colmap
# https://demuc.de/colmap/
def importColmapData(dataPath, useBlurred):

    def expandWavefrontGeometry_noUV(geometry):
        vertices = np.array(list(geometry.materials.items())[0][1].vertices)
        vertex_format = list(geometry.materials.items())[0][1].vertex_format
        if vertex_format == "V3F":
            vertices = vertices.reshape((int(vertices.shape[0] / 3), 3))
            vertices = np.concatenate((vertices[..., 0:3], np.zeros([vertices.shape[0], 2])), axis=1)  # XYZUV, set UV as zeros
        elif vertex_format == "N3F_V3F":
            vertices = vertices.reshape((int(vertices.shape[0] / 6), 6))
            vertices = np.concatenate((vertices[..., 3:6], np.zeros([vertices.shape[0], 2])), axis=1)  # XYZUV, set UV as zeros
        else:
            assert False, "vertex format of mesh does not match. It should either contain (normal+vertices) or (vertices)."
        mesh = EasyDict()
        mesh.vertices = vertices[..., :3].flatten()
        mesh.uvs = vertices[..., 3:].flatten()
        mesh.verticesAndUVs = vertices.flatten()
        mesh.faces = np.arange(int(vertices.shape[0])).flatten()
        return mesh

    modelPath = os.path.join(dataPath, "colmap", "dense", "0", "")  # including obj file
    registrationPath = os.path.join(dataPath, "colmap", "dense", "0", "images", "")  # including bundler file and images

    # parse colmap data
    parseColmapData(modelPath, registrationPath)

    # load mesh
    geometry = loadMesh(modelPath)
    mesh = expandWavefrontGeometry_noUV(geometry)

    # load texture (set it as Red if not found any)
    mesh.texture = loadTexture(modelPath, "colmap")

    # load undistorted images
    images = loadUndistortedImages(registrationPath, useBlurred)

    # load camera
    cameraSetups = loadCameraSetup(registrationPath, images)

    #modify Y and Z axes for mesh, camera, and viewMatrix
    trans_matrix1 = getAnisotropicScalingMatrix([1, -1, -1])  # Flip Y and Z for mesh and camera
    trans_matrix2 = getAnisotropicScalingMatrix([1, -1, -1])  # Flip Y and Z for viewMatrix
    transformMesh(mesh, trans_matrix1)
    for idx in range(len(cameraSetups)):
        cameraSetups[idx] = transformCamera(cameraSetups[idx], trans_matrix1)
        cameraSetups[idx].setViewMatrix(trans_matrix2.dot(cameraSetups[idx].viewMatrix))

    # ignore blacklist images
    images, cameraSetups = excludeBlacklist(registrationPath, images, cameraSetups)

    return mesh, cameraSetups, images

def loadMesh(modelPath):
    mesh_files = glob.glob(os.path.join(modelPath, "*.obj"))
    assert not len(mesh_files) == 0, "No obj mesh file found"
    assert not len(mesh_files) >= 2, "Only one obj Mesh file should be found in this folder"
    return pywavefront.Wavefront(mesh_files[0], collect_faces=True)

def loadTexture(modelPath, mode):
    if mode == "RealityCapture":
        textureFile = None
        for file in os.listdir(modelPath):
            if file.endswith("png"):
                textureFile = file
                break
        assert textureFile is not None, "RealityCapture data needs a texture file. No png texture file found!"
        return loadImage(modelPath + textureFile)
    elif mode == "colmap":
        texture = np.zeros([100, 100, 3], dtype=np.float32)  # make an arbitrary texture
        texture[:, :, 0] = 1  # set the texture as red
        return texture
    assert False, "mode is wrong for loading the texture!"

def loadCameraSetup(registrationPath, images):
    cameraFile = None
    for file in os.listdir(registrationPath):
        if file.endswith("out"):
            cameraFile = file
            break
    assert cameraFile is not None, "No bundler camera file found"

    cameraSetups = camerasFromBundler(
        registrationPath + cameraFile,
        [(img.shape[1], img.shape[0]) for img in images])

    return cameraSetups

def excludeBlacklist(registrationPath, images, cameraSetups):
    blacklistFile = registrationPath + 'blacklist.txt'
    blacklist = []
    if os.path.isfile(blacklistFile):
        blacklist = np.loadtxt(blacklistFile, dtype=int, ndmin=1)

    if len(blacklist) != 0:
        imgsTmp = []
        camsTmp = []
        for idx, (img, cam) in enumerate(zip(images, cameraSetups)):
            if idx not in blacklist:
                imgsTmp.append(img)
                camsTmp.append(cam)
        images = imgsTmp
        cameraSetups = camsTmp
    return images, cameraSetups

def loadUndistortedImages(registrationPath, useBlurred=False):
    # find the format of images in the directory
    if useBlurred is True:  # as the output of blur_background script is .png
        formatSuffix = ".png"
    else:
        formatSuffix = detectImageFormat(registrationPath)

    # load images
    imageFiles = []
    blurSuffix = "_blur"
    matteSuffix = "_matte"

    for file in sorted(os.listdir(registrationPath)):
        if useBlurred:
            if file.endswith(blurSuffix + formatSuffix):
                imageFiles.append(file)
        else:
            if file.endswith(formatSuffix) and not file.endswith(blurSuffix + formatSuffix) and not file.endswith(
                    matteSuffix + formatSuffix):
                imageFiles.append(file)
    assert len(imageFiles) != 0, "No image file found with suffix " + formatSuffix
    images = []
    for file in imageFiles:
        imagepath = registrationPath + file
        images.append(loadImage(imagepath))

    return images

#=======================================================

# export mesh using python pickle to speed up loading
def exportMeshAsPickle(mesh, filename):
    revert = mesh.gpu
    if revert:
        vao = mesh.vao
        textureGPU = mesh.textureGPU
        mesh.vao = None
        mesh.textureGPU = None
        mesh.gpu = False
    pickle.dump(mesh, open(filename, "wb") )
    if revert:
        mesh.vao = vao
        mesh.textureGPU = textureGPU
        mesh.gpu = True

#=======================================================

# export mesh as OBJ
def exportMeshAsObj(mesh, filename):
    hasUVs = mesh.uvs is not None
    verts = np.reshape(mesh.vertices, (int(mesh.vertices.shape[0] / 3), 3))
    if hasUVs:
        uvs = np.reshape(mesh.uvs, (int(mesh.uvs.shape[0] / 2), 2))
    faces = np.reshape(mesh.faces, (int(mesh.faces.shape[0] / 3), 3)) + 1
    with open(filename, 'w') as file:
        for v in verts:
            line = "v " + str(v[0]) + " " + str(v[1]) + " " + str(v[2])
            file.write(line + "\n")
        if hasUVs:
            for c in uvs:
                line = "vt " + str(c[0]) + " " + str(c[1])
                file.write(line + "\n")
        for f in faces:
            if hasUVs:
                a, b, c = [str(ff) for ff in f]
                line = "f " + a + "/" + a + " " +  b + "/" + b + " " + c + "/" + c
            else:
                line = "f " + str(f[0]) + " " + str(f[1]) + " " + str(f[2])
            file.write(line + "\n")
        file.close()